{ config, lib, ... }:
let
  cfg = config.khscodes.services.openssh;
in
{
  options.khscodes.services.openssh = {
    enable = lib.mkEnableOption "Enables openssh service for the instance";
    hostCertificate = {
      enable = lib.mkEnableOption "Enables getting host certificates from OpenBAO";
      path = lib.mkOption {
        type = lib.types.str;
        default = "ssh-host";
      };
      hostNames = lib.mkOption {
        type = lib.types.listOf lib.types.str;
        description = "The list of host names to get certificates for";
        default = [ ];
      };
    };
  };

  config = lib.mkIf cfg.enable (
    let
      certificateNames = lib.lists.unique cfg.hostCertificate.hostNames;
      hostCertificatEnable = cfg.hostCertificate.enable && cfg.hostCertificate.hostNames != [ ];
      vaultRoleName = config.khscodes.infrastructure.vault-server-approle.role_name;
      fqdn = config.networking.fqdnOrHostName;
      sshHostBackend = "ssh-host";
    in
    {
      services.openssh = {
        enable = true;
        settings = {
          PasswordAuthentication = false;
          PermitRootLogin = "no";
          KbdInteractiveAuthentication = false;
        };
        extraConfig = lib.mkIf hostCertificatEnable ''
          HostCertificate /etc/ssh/ssh_host_ed25519_key-cert.pub
        '';
      };
      khscodes.infrastructure.vault-server-approle = {
        enable = true;
        policy."${sshHostBackend}/sign/${vaultRoleName}" = {
          capabilities = [
            "read"
            "update"
            "create"
          ];
        };
        stageModules = [
          {
            khscodes.vault.ssh_secret_backend_role.${vaultRoleName} = {
              name = fqdn;
              backend = cfg.hostCertificate.path;
              key_type = "ca";
              allow_host_certificates = true;
              allow_bare_domains = true;
              allowed_domains = certificateNames;
              allowed_user_key_config = [
                {
                  type = "ed25519";
                  lengths = [ 0 ];
                }
              ];
            };
          }
        ];
      };
      khscodes.services.vault-agent = lib.mkIf hostCertificatEnable {
        enable = true;
        templates = [
          {
            contents = ''
              {{- $public_key := file "/etc/ssh/ssh_host_ed25519_key.pub" -}}
              {{- $public_key = printf "public_key=%s" $public_key -}}
              {{- with secret "ssh-host/sign/${fqdn}" "cert_type=host" $public_key "valid_principals=${lib.strings.concatStringsSep "," certificateNames}" -}}
              {{ .Data.signed_key }}
              {{- end -}}
            '';
            destination = "/etc/ssh/ssh_host_ed25519_key-cert.pub";
            perms = "0644";
            restartUnits = [ "sshd.service" ];
          }
        ];
      };
    }
  );
}
