{ config, lib, pkgs, ... }:

# TODO: should max-builds be enforced on thisHost as well?

let
  inherit (builtins)
    toString
    elem
    attrNames
    attrValues
    ;

  # TODO: test ssh-ng
  # https://discourse.nixos.org/t/wrapper-to-restrict-builder-access-through-ssh-worth-upstreaming/25834
  nix-ssh-wrapper = pkgs.writeShellScript "nix-ssh-wrapper" ''
    case $SSH_ORIGINAL_COMMAND in
      "nix-daemon --stdio")
        exec ${config.nix.package}/bin/nix-daemon --stdio
        ;;
      "nix-store --serve --write")
        exec ${config.nix.package}/bin/nix-store --serve --write
        ;;
      *)
        echo "Access only allowed for using the nix remote builder" 1>&2
        exit
    esac
  '';

  known-hosts = let
    known-hosts' = lib.importTOML ./hosts.toml; # TODO: eww
  in
    lib.pipe known-hosts' [
      (lib.flip lib.removeAttrs ["__default__"])
      (lib.mapAttrs (fqdn: host:
        lib.recursiveUpdate (known-hosts'."__default__" or {}) host
      ))
      (lib.mapAttrsToList (fqdn: host:
        let
          allHostNames = [ fqdn ] ++ host.aliases;
        in
        lib.forEach allHostNames (alias:
          lib.nameValuePair
            alias
            (host // {
              aliases = lib.remove alias allHostNames;
              isAlias = fqdn != alias;
            })
        )
      ))
      lib.flatten
      lib.listToAttrs
    ];
  hostNames = attrNames known-hosts;
  thisHost = known-hosts.${config.networking.fqdn} or known-hosts.${config.networking.hostName};
  thisHostIsBuilder = thisHost.buildMachine.maxJobs > 0 && thisHost.ssh ? listenPublicKey;
  thisHostIsBuildee = thisHost.ssh ? userPublicKey;
  thisHostIsHopHost = elem config.networking.fqdn (lib.forEach (attrValues known-hosts) (host: host.ssh.proxyJump or null));

  mkRemoteConfig = fqdn: let
    thatName = lib.head lib.splitString "." fqdn;
    thatDomain = lib.concatStringsSep "." (lib.tail (lib.splitString "." fqdn));

    thatHost = known-hosts.${fqdn};
    thatJump = known-hosts.${thatHost.ssh.proxyJump};
    buildMachine = thatHost.buildMachine // {
      hostName = fqdn;
      sshUser = thatHost.ssh.listenUser;
    };
    remoteStore = "${buildMachine.protocol}://${buildMachine.sshUser}@${buildMachine.hostName}";

    thatHostIsBuilder = thatHost.buildMachine.maxJobs > 0 && thatHost.ssh ? listenPublicKey;
    thatHostIsBuildee = thatHost.ssh ? userPublicKey && thisHostIsBuilder;
    thatHostIsThis = elem config.networking.fqdn ([ fqdn ] ++ thatHost.aliases);
  in lib.mkIf (!thatHostIsThis) ( lib.mkMerge [
    # out
    (lib.mkIf (thisHostIsBuildee && thatHostIsBuilder) {

      # TODO: Allow setting speedFactor for local builds, as local is currently fixed to 0
      # https://github.com/NixOS/nix/issues/2457

      nix.distributedBuilds = true;
      # useful when the builder has a faster internet connection than i do
      nix.settings.builders-use-substitutes = true;
      nix.buildMachines = lib.mkIf (!thatHost.isAlias) [ buildMachine ];

      nix.settings.substituters = lib.mkIf (thatHost.useAsSubstituter && config.currentSpecialisation != "remote-store-${fqdn}") [
        "${remoteStore}?trusted=true"
      ];

      specialisation = lib.mkIf (thatHost.remoteStoreSpecialization or false && !thatHost.isAlias) {
        "remote-store-${fqdn}" = {
          inheritParentConfig = true;
          configuration = {
            currentSpecialisation = lib.mkOverride 0 "remote-store-${fqdn}";
            # https://docs.lix.systems/manual/lix/stable/command-ref/conf-file.html#conf-store
            # https://nix.dev/manual/nix/stable/command-ref/conf-file.html#conf-store
            # https://nix.dev/manual/nix/stable/store/types/
            nix.settings.store = "${remoteStore}?trusted=true";
          };
        };
      };

    })
    # out or jump
    (lib.mkIf (thisHostIsBuildee && thatHost.ssh ? listenPublicKey) {
      programs.ssh.knownHosts.${fqdn}.publicKey = thatHost.ssh.listenPublicKey;
      # TODO: use nix.buildMachines.*.publicHostKey ?

      # timeouts are great when remote is unresponsive. nix doesn't care, lix is way and tests each remote only once
      programs.ssh.extraConfig = ''
        Host ${fqdn}
          Port ${toString thatHost.ssh.listenPort}
          ${lib.optionalString ((thatHost.ssh.connectTimeout or 0) != 0) ''
            ConnectTimeout ${toString thatHost.ssh.connectTimeout}
          ''}
          ${lib.optionalString (thatHost.ssh ? proxyJump && !lib.elem thatDomain (thatHost.noProxyJumpDomains or [])) ''
            ProxyJump ${thatJump.ssh.listenUser}@${thatHost.ssh.proxyJump}:${toString thatJump.ssh.listenPort}
          ''}
          ${lib.optionalString (thatHost.ssh ? userPrivateKey) ''
            IdentityFile ${thatHost.ssh.userPrivateKey}
          ''}
      '';

      sops.secrets = lib.mkIf (lib.hasPrefix "/run/secrets/" (thatHost.ssh.userPrivateKey or "")) {
        "${lib.removePrefix "/run/secrets/" thatHost.ssh.userPrivateKey}" = { };
      };

    })
    # in
    (lib.mkIf ((thisHostIsBuilder || thisHostIsHopHost) && thatHostIsBuildee && !thatHost.isAlias) {
      # TODO: ensure the user is "nixbld-remote"?
      users.groups.${thisHost.ssh.listenUser} = { };
      users.users.${thisHost.ssh.listenUser} = {
        isSystemUser = lib.mkDefault (!config.users.users.${thisHost.ssh.listenUser}.isNormalUser);
        useDefaultShell = lib.mkDefault true;
        openssh.authorizedKeys.keys = [
          # https://man.archlinux.org/man/core/openssh/sshd.8.en#AUTHORIZED_KEYS_FILE_FORMAT
          # TODO: lib.getExe
          ''restrict,pty,command="${nix-ssh-wrapper}" ${thatHost.ssh.userPublicKey}''
        ];
        group = lib.mkOverride 1499 "${thisHost.ssh.listenUser}"; # mkOptionDefault - 1,
      };
      # TODO: https://nix.dev/tutorials/nixos/distributed-builds-setup.html#optimise-the-remote-builder-configuration
    })
    (lib.mkIf (thisHostIsBuilder && thatHostIsBuildee && !thatHost.isAlias) {
      nix.settings.allowed-users = [ thisHost.ssh.listenUser ];
      nix.settings.trusted-users = [ thisHost.ssh.listenUser ];
    })
  ]);

in {

  imports = lib.forEach hostNames mkRemoteConfig;

  # TODO: upstream this as specialisation.currentSpecialization that is `nullOr str`
  # https://github.com/NixOS/nixpkgs/blob/b6eaf97c6960d97350c584de1b6dcff03c9daf42/nixos/modules/system/activation/specialisation.nix#L77
  # https://github.com/NixOS/nixpkgs/blob/b6eaf97c6960d97350c584de1b6dcff03c9daf42/nixos/modules/system/activation/no-clone.nix
  options.currentSpecialisation = lib.mkOption {
    type = lib.types.nullOr lib.types.str;
    internal = true;
    default = null;
    description = "Which specialization this is, if any.";
  };

  config = {
    nix.settings.max-jobs =
      lib.mkIf ((thisHost.buildMachine.maxJobs or 0) > 0)
      (lib.mkDefault thisHost.buildMachine.maxJobs);
  };

  # TODO:
  # nix.buildMachines = [
  #   {
  #     protocol = null;
  #     hostName = "localhost";
  #     speedFactor = SOMETHING; // perhaps by looking up self in hosts.toml, perhaps by always preferring self, perhaps by always preferring remotes
  #     maxJobs = config.nix.settings.SOMETHING;
  #     system = config.nix.settings.SOMETHING;
  #     supportedFeatures = config.nix.settings.SOMETHING;
  #   }
  # ];


}