{ config, lib, ... }:

with lib;

let
  cfg = config.mine.ssh-on-boot;
in {
  options.mine.ssh-on-boot = {
    enable = mkOption {
      type = types.bool;
      default = false;
    };
    
    network = {
      address = mkOption {
        type = types.str;
        example = "192.168.1.11";
      };

      gateway = mkOption {
        type = types.str;
        example = "192.168.1.1";
      };

      netmask = mkOption {
        type = types.str;
        example = "255.255.255.0";
      };

      hostname = mkOption {
        type = types.str;
        default = "${config.networking.hostName}-boot";
      };

      interface = mkOption {
        type = types.str;
        example = "eno3";
      };
    };

    kernelModules = mkOption {
      type = types.listOf types.str;
      default = [
        "ixgbe"
        "igb"
      ];
    };

    sshPort = mkOption {
      type = types.int;
      default = 2222;
    };

    sshKeyLocation = mkOption {
      type = types.str;
      default = "/state/root/ssh-on-boot";
    };
  };

  config = mkIf cfg.enable {
    boot = {
      kernelParams = [
        "ip=${cfg.network.address}::${cfg.network.gateway}:${cfg.network.netmask}:${cfg.network.hostname}:${cfg.network.interface}"
      ];

      initrd.availableKernelModules = cfg.kernelModules;
      initrd.network = {
        enable = true;
  
        ssh = {
          enable = true;
          port = cfg.sshPort;
          hostKeys = [
            "${cfg.sshKeyLocation}/ssh_host_ed25519_key"
            "${cfg.sshKeyLocation}/ssh_host_rsa_key"
          ];

          authorizedKeys = config.users.users.root.openssh.authorizedKeys.keys;
        };
  
        postCommands = let
          luksCmd = builtins.concatStringsSep "; " (
            lib.mapAttrsToList (n: v:
              "echo Opening ${n}; cryptsetup-askpass open ${v.device} ${n}"
            ) config.boot.initrd.luks.devices);
        in ''
          ip route add ${cfg.network.gateway} dev ${cfg.network.interface}
          ip route add default via ${cfg.network.gateway} dev ${cfg.network.interface}
          ip link set ${cfg.network.interface} up

          echo "${luksCmd}; zpool import -a; zfs load-key -a; killall zfs" >> /root/.profile
        '';
      };
    };
  };
}