set -euo pipefail

usage() {
  cat <<'USAGE'
Usage:
  nix run .#bootstrap-host -- <main-wiki|replica-wiki> <target-host> [ssh-identity-file]
  nix run .#bootstrap-host -- <main-target-host> <replica-target-host> [ssh-identity-file]

Examples:
  nix run .#bootstrap-host -- main-wiki root@203.0.113.10 ~/.ssh/do-bootstrap
  nix run .#bootstrap-host -- root@203.0.113.10 root@203.0.113.11 ~/.ssh/do-bootstrap
USAGE
}

if [ "$#" -lt 2 ] || [ "$#" -gt 3 ]; then
  usage
  exit 1
fi

repo_root="$(pwd)"
if [ ! -f "$repo_root/flake.nix" ]; then
  printf 'Run bootstrap-host from the repo root\n' >&2
  exit 1
fi

ssh_identity_file=""
main_target=""
replica_target=""

case "$1" in
  main-wiki|replica-wiki)
    if [ "$1" = "main-wiki" ]; then
      main_target="$2"
    else
      replica_target="$2"
    fi
    ssh_identity_file="${3:-}"
    ;;
  *)
    main_target="$1"
    replica_target="$2"
    ssh_identity_file="${3:-}"
    ;;
esac

make_host_module() {
  local module_file="$1"

  cat > "$module_file" <<'MODULE'
{ ... }:
{
  services.journald.storage = "persistent";

  services.openssh = {
    enable = true;
    openFirewall = true;
    settings = {
      AllowAgentForwarding = false;
      AllowGroups = [ "wheel" ];
      AllowTcpForwarding = false;
      ClientAliveCountMax = 2;
      ClientAliveInterval = 300;
      PasswordAuthentication = false;
      KbdInteractiveAuthentication = false;
      LoginGraceTime = 20;
      MaxAuthTries = 3;
      MaxSessions = 4;
      PermitRootLogin = "no";
      PermitTunnel = false;
      PermitUserEnvironment = false;
      StreamLocalBindUnlink = false;
      X11Forwarding = false;
    };
  };

  users.users.jet = {
    isNormalUser = true;
    extraGroups = [ "wheel" ];
    openssh.authorizedKeys.keys = [
@ADMIN_KEYS@
    ];
  };

  security.sudo.wheelNeedsPassword = false;

  services.do-agent.enable = false;
}
MODULE
}

run_bootstrap() {
  local host_name="$1"
  local target_host="$2"
  local work_dir
  local module_file
  local remote_target
  local try
  local ssh_cmd=(ssh -o StrictHostKeyChecking=accept-new)
  local scp_cmd=(scp -o StrictHostKeyChecking=accept-new)

  work_dir="$(mktemp -d)"
  module_file="$work_dir/host-bootstrap.nix"
  remote_target="$target_host:/etc/nixos/host-bootstrap.nix"

  make_host_module "$module_file"

  if [ -n "$ssh_identity_file" ]; then
    ssh_cmd+=( -i "$ssh_identity_file" )
    scp_cmd+=( -i "$ssh_identity_file" )
  fi

  "${ssh_cmd[@]}" "$target_host" 'mkdir -p /etc/nixos'
  "${scp_cmd[@]}" "$module_file" "$remote_target"

  printf 'Infecting %s onto %s\n' "$host_name" "$target_host"
  "${ssh_cmd[@]}" "$target_host" \
    'umount /boot/efi 2>/dev/null || true; curl -fsSL https://raw.githubusercontent.com/elitak/nixos-infect/36f48d8feb89ca508261d7390355144fc0048932/nixos-infect | env PROVIDER=digitalocean doNetConf=y NIX_CHANNEL=nixos-24.05 NIXOS_IMPORT=./host-bootstrap.nix bash -x' || true

  printf 'Waiting for %s to reboot into NixOS\n' "$host_name"
  for try in $(seq 1 60); do
    if "${ssh_cmd[@]}" -o ConnectTimeout=5 "$target_host" 'grep -q "^ID=nixos" /etc/os-release'; then
      break
    fi
    sleep 5
  done

  printf 'Finalizing network config on %s\n' "$host_name"
  "${ssh_cmd[@]}" "$target_host" '
    sed -i "/defaultGateway6 = {/,/};/d" /etc/nixos/networking.nix 2>/dev/null || true
    sed -i "/ipv6.routes = \[ { address = \"\"; prefixLength = 128; } \];/d" /etc/nixos/networking.nix 2>/dev/null || true
    nixos-rebuild switch
  '

  rm -rf "$work_dir"
}

if [ -n "${main_target:-}" ]; then
  run_bootstrap main-wiki "$main_target"
fi
if [ -n "${replica_target:-}" ]; then
  run_bootstrap replica-wiki "$replica_target"
fi

printf '\nBootstrap complete. The hosts should now be reachable as minimal NixOS systems over public SSH.\n'