#!/usr/bin/env bash
# recover-hosts.sh — Fix sops/SSH-key/GitHub-token issues on deployed NixOS hosts
# and trigger a Switch-nix rebuild on each.
#
# Run from the repo root on the workstation (nixos@nixos):
#   bash scripts/recover-hosts.sh [<hostname> ...]
#
# With no args it discovers and checks every known hostname.
# With args it checks only those hostnames:
#   bash scripts/recover-hosts.sh tor-relay
#
# Fixes applied automatically (then prompts before rebuilding):
#   1. SSH host key drift — live key no longer matches host-keys/<target>_ssh_host_ed25519_key
#      Fix: scp the registered key back and restore it (needs sudo once per host).
#      To push new keys proactively (before drift, e.g. right after
#      sync-host-keys.sh --regenerate-all-keys), use instead:
#        scripts/secrets/push-host-keys.sh --all
#   2. Stale/invalid GitHub access token — the rendered nix-github-token.conf has
#      a token GitHub rejects (401), blocking any rebuild that fetches disko or
#      other public GitHub flake inputs.
#      Fix: empty the rendered file so nix makes unauthenticated requests instead.
#      Public repos (disko, nixpkgs, etc.) work fine without auth. sops-nix
#      re-renders the correct new token automatically after the first successful
#      rebuild.
#
# Both fixes need one interactive sudo session per host. The script opens a
# single ssh -t per broken host so you enter the password once and all steps
# run in sequence.

set -euo pipefail
cd "$(dirname "$0")/.."
source scripts/env.sh 2>/dev/null || true

SSH_OPTS=(-o StrictHostKeyChecking=no -o BatchMode=yes -o ConnectTimeout=5)
SSH_USER=nixos

# Known flake-target → ssh hostname map for all currently-defined hosts.
# Add new hosts here as they are deployed.
declare -A TARGET_HOST=(
  [lxc-docker]=docker
  [lxc-nix-cache]=nix-cache
  [lxc-pxe-boot]=pxe-boot
  [lxc-tor-relay]=tor-relay
  [lxc-minimal]=nix-minimal
  [proxmox-server]=server
  [baremetal-gui]=nixos
)

# ── helpers ───────────────────────────────────────────────────────────────────

info() { echo "  [✓] $*"; }
warn() { echo "  [!] $*"; }
step() { echo "==> $*"; }

ssh_host_age() {
  ssh-keyscan -t ed25519 "$1" 2>/dev/null \
    | nix shell nixpkgs#ssh-to-age --command ssh-to-age 2>/dev/null \
    | head -1 || true
}

registered_age() {
  local keyfile="host-keys/${1}_ssh_host_ed25519_key.pub"
  [ -f "$keyfile" ] || return 0
  nix shell nixpkgs#ssh-to-age --command ssh-to-age < "$keyfile" 2>/dev/null \
    | head -1 || true
}

github_token_valid() {
  local host=$1
  local raw token code
  raw=$(ssh "${SSH_OPTS[@]}" "$SSH_USER@$host" \
    "cat /run/secrets/rendered/nix-github-token.conf 2>/dev/null || true")
  token=$(echo "$raw" | grep -oP '(?<=github\.com=)\S+' || true)
  if [ -z "$token" ]; then
    return 0  # no token = unauthenticated, works for public repos
  fi
  code=$(curl -s -o /dev/null -w "%{http_code}" \
    -H "Authorization: token $token" \
    "https://api.github.com/repos/nix-community/disko" 2>/dev/null || echo 000)
  [ "$code" = "200" ]
}

# ── discover hosts ────────────────────────────────────────────────────────────

if [ $# -gt 0 ]; then
  HOSTNAMES=("$@")
else
  HOSTNAMES=()
  seen=()
  for target in "${!TARGET_HOST[@]}"; do
    h="${TARGET_HOST[$target]}"
    # deduplicate (e.g. proxmox-server and lxc-server both map to "server")
    if [[ ! " ${seen[*]:-} " =~ " $h " ]]; then
      seen+=("$h")
      if ssh "${SSH_OPTS[@]}" "$SSH_USER@$h" "true" 2>/dev/null; then
        HOSTNAMES+=("$h")
      fi
    fi
  done
fi

if [ ${#HOSTNAMES[@]} -eq 0 ]; then
  echo "No reachable hosts found. Pass hostnames explicitly or check SSH."
  exit 1
fi

echo ""
echo "Hosts to check: ${HOSTNAMES[*]}"
echo ""

# ── check phase ───────────────────────────────────────────────────────────────

NEEDS_FIX=()

for host in "${HOSTNAMES[@]}"; do
  step "$host"

  if ! ssh "${SSH_OPTS[@]}" "$SSH_USER@$host" "true" 2>/dev/null; then
    warn "SSH unreachable — clearing stale known_hosts entry"
    ssh-keygen -R "$host" 2>/dev/null || true
    continue
  fi

  flake_target=$(ssh "${SSH_OPTS[@]}" "$SSH_USER@$host" \
    "cat /etc/flake-target 2>/dev/null || true")
  echo "  flake-target: ${flake_target:-unknown}"

  host_broken=false

  # SSH host key
  if [ -n "$flake_target" ] && [ -f "host-keys/${flake_target}_ssh_host_ed25519_key.pub" ]; then
    live=$(ssh_host_age "$host")
    want=$(registered_age "$flake_target")
    if [ "$live" = "$want" ]; then
      info "SSH host key OK"
    else
      warn "SSH host key MISMATCH (live ≠ host-keys/) -- use push-host-keys.sh proactively next time"
      echo "       live:       $live"
      echo "       registered: $want"
      host_broken=true
    fi
  else
    echo "  [~] No host-keys/ entry for ${flake_target:-unknown} — skipping key check"
  fi

  # GitHub token
  if github_token_valid "$host"; then
    info "GitHub token OK"
  else
    warn "GitHub token invalid (rebuild will fail with 401)"
    host_broken=true
  fi

  # sops-nix result
  sops_result=$(ssh "${SSH_OPTS[@]}" "$SSH_USER@$host" \
    "systemctl show sops-nix --property=Result --value 2>/dev/null || echo unknown")
  if [ "$sops_result" = "success" ]; then
    info "sops-nix: success"
  else
    warn "sops-nix: $sops_result"
  fi

  $host_broken && NEEDS_FIX+=("$host")
  echo ""
done

# ── fix phase ─────────────────────────────────────────────────────────────────

if [ ${#NEEDS_FIX[@]} -eq 0 ]; then
  echo "All hosts healthy — nothing to fix."
  exit 0
fi

echo "Hosts needing fixes: ${NEEDS_FIX[*]}"
echo ""
echo "Each fix requires one sudo session per host. You will be prompted for"
echo "the nixos sudo password once per host; all steps run in that session."
echo ""
read -r -p "Proceed with fixes + Switch-nix on each broken host? [y/N] " confirm
[[ "$confirm" =~ ^[Yy]$ ]] || { echo "Aborted."; exit 0; }
echo ""

for host in "${NEEDS_FIX[@]}"; do
  step "Fixing $host"
  flake_target=$(ssh "${SSH_OPTS[@]}" "$SSH_USER@$host" \
    "cat /etc/flake-target 2>/dev/null || true")

  fix_script=""

  # Fix 1: restore SSH host key
  live=$(ssh_host_age "$host")
  want=$(registered_age "${flake_target:-}")
  if [ -n "$want" ] && [ "$live" != "$want" ]; then
    echo "  Uploading registered SSH host key (private + public)..."
    scp -o StrictHostKeyChecking=no \
      "host-keys/${flake_target}_ssh_host_ed25519_key" \
      "$SSH_USER@$host:/tmp/recover_ed25519_key"
    scp -o StrictHostKeyChecking=no \
      "host-keys/${flake_target}_ssh_host_ed25519_key.pub" \
      "$SSH_USER@$host:/tmp/recover_ed25519_key.pub"
    fix_script+='
echo "[fix] Restoring SSH host key..."
install -m 0600 /tmp/recover_ed25519_key /etc/ssh/ssh_host_ed25519_key
install -m 0644 /tmp/recover_ed25519_key.pub /etc/ssh/ssh_host_ed25519_key.pub
rm -f /tmp/recover_ed25519_key /tmp/recover_ed25519_key.pub
echo "      Done."
'
    ssh-keygen -R "$host" 2>/dev/null || true
  fi

  # Fix 2: clear invalid GitHub token
  if ! github_token_valid "$host"; then
    fix_script+='
echo "[fix] Clearing stale GitHub token (nix will use unauthenticated access)..."
echo "" > /run/secrets/rendered/nix-github-token.conf
systemctl restart nix-daemon 2>/dev/null || true
echo "      Done."
'
  fi

  # Fix 3: rebuild
  fix_script+='
echo "[fix] Running nixos-rebuild switch..."
nixos-rebuild switch \
  --no-write-lock-file \
  --refresh \
  --flake "git+https://gitea.lan.ddnsgeek.com/beatzaplenty/nixos.git#$(cat /etc/flake-target)"
echo "[fix] Rebuild complete."
'

  echo "  Opening SSH session (enter sudo password when prompted)..."
  if ssh -t -o StrictHostKeyChecking=no "$SSH_USER@$host" \
    "sudo bash -s" <<< "$fix_script"; then
    echo ""
    info "$host fixed and rebuilt"
  else
    rc=$?
    echo ""
    warn "$host: rebuild exited with code $rc (may still have succeeded — check sops-nix below)"
  fi

  # Verify: re-check sops-nix result post-rebuild
  sops_result_after=$(ssh "${SSH_OPTS[@]}" "$SSH_USER@$host" \
    "systemctl show sops-nix --property=Result --value 2>/dev/null || echo unknown" 2>/dev/null || echo "ssh-failed")
  if [ "$sops_result_after" = "success" ]; then
    info "$host sops-nix: success post-rebuild"
  else
    warn "$host sops-nix: $sops_result_after post-rebuild (may need another pass)"
  fi
  echo ""
done

echo "Recovery complete."