# New attempt at consolidating SSH config that wraps the step client, SSH certs, and all that in a top-level module
{ self, inputs, lib, config, ... }:
let
  sshHostCAPath = ../../hosts/janus/public/ssh_host_ca_key.pub;
  caPatterns = [ "*.john-stream.com" "192.168.1.*" "fded:fb16:653e:25da:be24:11ff:*" ];
  wrappers = inputs.self.wrappers;
  sshCertConfig = config.optionModules.ssh-certs;

  mkHostScripts = { cfg, pkgs, provisionerPasswordFile ? null }: let
    hostKeyFile = "${cfg.host.configDir}/${cfg.host.keyFile}";
    hostCertFile = "${hostKeyFile}-cert.pub";
    wrapperCfg = cfg // {
      certificates = cfg.certificates // {
        provisionerPasswordFile = provisionerPasswordFile;
        overwrite = true;
      };
    };
  in {
    inherit hostKeyFile hostCertFile;
    sign = (wrappers.signHostWrapper.apply {
      inherit pkgs;
      ssh-new = wrapperCfg;
    }).wrapper;
    renew = (wrappers.renewHostWrapper.apply {
      inherit pkgs;
      ssh-new = wrapperCfg;
    }).wrapper;
    check = (wrappers.hostCheckWrapper.apply {
      inherit pkgs;
      ssh-new = wrapperCfg;
    }).wrapper;
    renewalCheck = (wrappers.renewalCheck.apply {
      inherit pkgs;
      ssh-new = wrapperCfg;
    }).wrapper;
  };

  mkUserScripts = { cfg, pkgs, provisionerPasswordFile ? null }: {
    sign = let
      wrapperCfg = cfg // {
        certificates = cfg.certificates // {
          provisionerPasswordFile = provisionerPasswordFile;
          overwrite = true;
        };
      };
    in (wrappers.signUserWrapper.apply {
      inherit pkgs;
      ssh-new = wrapperCfg;
    }).wrapper;
    check = (wrappers.userCheckWrapper.apply {
      inherit pkgs;
      ssh-new = cfg;
    }).wrapper;
  };
in
{
  flake.modules.nixos.ssh-new = { config, pkgs, lib, ... }:
    let
      cfg = config.ssh-new;
      hasAdminJwk = lib.hasAttrByPath [ "sops" "secrets" "janus/admin_jwk" ] config;
      adminJwkPath = if hasAdminJwk then config.sops.secrets."janus/admin_jwk".path else null;
      hostScripts = mkHostScripts {
        inherit cfg pkgs;
        provisionerPasswordFile = if cfg.certificates.host.autoRenew
          then adminJwkPath
          else null;
      };
      userScripts = mkUserScripts {
        inherit cfg pkgs;
        provisionerPasswordFile = adminJwkPath;
      };
      CAknownHosts = (lib.genAttrs caPatterns (_: {
        certAuthority = true;
        publicKey = lib.removeSuffix "\n" (builtins.readFile sshHostCAPath);
      }));
    in
    {
      imports = [ sshCertConfig ];
      config = {
        services.openssh = {
          enable = true;
          openFirewall = true;
          hostKeys = [
            {
              path = hostScripts.hostKeyFile;
              type = cfg.host.keyType;
            }
          ];
          settings = lib.mkMerge [
            {
              PasswordAuthentication = false;
              KbdInteractiveAuthentication = false;
            }
            cfg.host.extraSettings
            (lib.mkIf cfg.certificates.host.enable {
              HostCertificate = hostScripts.hostCertFile;
            })
            (lib.mkIf cfg.certificates.user.enable {
              TrustedUserCAKeys = "${cfg.host.configDir}/${cfg.certificates.user.CAFile}";
            })
          ];
        };

        environment.etc."ssh/${cfg.certificates.user.CAFile}" = lib.mkIf cfg.certificates.user.enable {
          source = ../../hosts/janus/public/ssh_user_ca_key.pub;
        };

        environment.systemPackages = (
          lib.optionals cfg.certificates.host.enable [
            hostScripts.sign
            hostScripts.renew
            hostScripts.renewalCheck
            hostScripts.check
          ]
        );

        sops = lib.mkIf cfg.certificates.host.autoRenew {
          secrets."janus/admin_jwk" = {
            sopsFile = ../../../keys/secrets.yaml;
            owner = "root";
            group = "root";
            mode = "0400";
          };
        };

        systemd = lib.mkIf cfg.certificates.host.autoRenew {
          services.ssh-certs-renew = {
            description = "SSH host certificate renewal";
            wantedBy = [ "multi-user.target" ];
            after = [ "network-online.target" ];
            wants = [ "network-online.target" ];
            path = with pkgs; [ step-cli systemd ];
            serviceConfig = {
              Type = "oneshot";
              User = "root";
              Group = "root";
              ExecCondition = lib.getExe hostScripts.renewalCheck;
              ExecStart = lib.getExe hostScripts.renew;
            };
          };
          timers.ssh-certs-renew = {
            description = "Periodic Step SSH host certificate renewal";
            wantedBy = [ "timers.target" ];
            timerConfig = {
              OnBootSec = "5m";
              OnUnitActiveSec = "4h";
              RandomizedDelaySec = "15m";
              Persistent = true;
              Unit = "ssh-certs-renew.service";
            };
          };
        };

        # This winds up in /etc/ssh/ssh_known_hosts, which applies to all users
        programs.ssh.knownHosts = lib.mkIf cfg.certificates.user.enable CAknownHosts;

        home-manager.users.root = lib.mkIf cfg.certificates.user.enable {
          home.stateVersion = lib.mkDefault config.system.stateVersion;
          imports = [ inputs.self.modules.homeManager.ssh-new ];
          home.packages = [
            userScripts.sign
            userScripts.check
          ];
          ssh-new.certificates.user.enable = true;
        };
      };
    };

  flake.modules.homeManager.ssh-new = { config, pkgs, lib, ... }:
    let
      cfg = config.ssh-new;
      hasAdminJwk = lib.hasAttrByPath [ "sops" "secrets" "janus/admin_jwk" ] config;
      adminJwkPath = if hasAdminJwk then config.sops.secrets."janus/admin_jwk".path else null;
      hostScripts = mkHostScripts { inherit cfg pkgs; };
      sshHostCAContent = lib.removeSuffix "\n" (builtins.readFile sshHostCAPath);
      knownHostsText = lib.concatMapStrings
        (pattern: "@cert-authority ${pattern} ${sshHostCAContent}\n")
        caPatterns;
      userScripts = mkUserScripts {
        inherit cfg pkgs;
        provisionerPasswordFile = adminJwkPath;
      };
    in
    {
      imports = [ sshCertConfig ];
      config = {
        home.file.".ssh/known_hosts" = lib.mkIf cfg.certificates.user.enable {
          text = knownHostsText;
        };
        home.packages = [
          userScripts.sign
          userScripts.check
        ]
        ++ lib.optionals cfg.host.enable-scripts [
          hostScripts.sign
          hostScripts.renew
          hostScripts.renewalCheck
          hostScripts.check
        ];
        programs.ssh = {
          enable = true;
          enableDefaultConfig = false;
          settings = {
            # These settings apply to all connections
            "*" = lib.mkMerge (
              [
                # Default settings
                {
                  Compression = false;
                  ServerAliveInterval = 60;
                  ServerAliveCountMax = 3;
                  TCPKeepAlive = "yes";
                  ConnectTimeout = 3;

                  PubkeyAuthentication = "yes";
                  PasswordAuthentication = "no";
                  PreferredAuthentications = "publickey";
                  IdentitiesOnly = true;
                  IdentityFile = cfg.user.keyFile;

                  StrictHostKeyChecking = "accept-new";
                  UserKnownHostsFile = "${config.home.homeDirectory}/.ssh/known_hosts";

                  HashKnownHosts = lib.mkDefault false;
                  AddKeysToAgent = lib.mkDefault "yes";
                  ForwardAgent = lib.mkDefault false;

                  RequestTTY = lib.mkDefault "auto";
                  SetEnv.TERM = "xterm-256color";
                }
              ]
              # SSH certificate settings
              ++ lib.optionals cfg.certificates.user.enable [
                {
                  CertificateFile = "${cfg.user.keyFile}-cert.pub";
                }
              ]
            );
            "john-pc" = {
              HostName = "192.168.1.85";
              User = "john";
            };
            "gitea" = {
              HostName = "192.168.1.104";
              User = "john";
            };
            "janus" = {
              HostName = "fded:fb16:653e:25da:be24:11ff:fe6b:4d57";
              User = "root";
            };
            "soteria" = {
              HostName = "fded:fb16:653e:25da:be24:11ff:fe54:aa39";
              User = "root";
            };
            "hermes" = {
              HostName = "192.168.1.150";
              User = "root";
            };
            "jdl-docker" = {
              HostName = "jdl-docker.tailcf205.ts.net";
              User = "john";
            };
          };
        };
      };
    };
}