#!/usr/bin/env bash
# gc-hosts.sh — Run nix-collect-garbage -d on all live NixOS hosts and pve1.
#
# nix-cache is excluded: it is the shared binary cache for all other hosts, so
# gc-ing it would evict cached store paths and force costly rebuilds elsewhere.
#
# Runs SSH jobs in parallel (up to MAX_JOBS at a time) and prints a summary.
#
# Usage (from repo root):
#   bash scripts/gc-hosts.sh [--dry-run] [<hostname> ...]
#
# Options:
#   --dry-run     Print the SSH commands without executing them.
#   <hostname>    Limit to the given hostnames (bare names, no domain suffix).
#                 Default: all hosts in the list below.
#
# Sudo notes:
#   NixOS hosts: tries "sudo -n nix-collect-garbage -d" first (non-interactive,
#   works when wheelNeedsPassword = false such as on the HA cluster).  Falls
#   back to "nix-collect-garbage -d" as the nixos user if sudo requires a
#   password — this still collects unreferenced store paths and removes old
#   nixos-user profile generations, but will not remove old system generations.
#   pve1: non-NixOS Proxmox node, runs nix-collect-garbage -d as wayne (no
#   system generations to delete).

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

# ── config ────────────────────────────────────────────────────────────────────

: "${MAX_JOBS:=6}"
: "${NIXOS_USER:=nixos}"
SSH_OPTS=(-o StrictHostKeyChecking=no -o BatchMode=yes -o ConnectTimeout=10)

# Map: hostname → ssh-user
# Update this list when new hosts are added/removed.
declare -A HOSTS=(
  [docker]="$NIXOS_USER"
  [ha-server-1]="$NIXOS_USER"
  [ha-server-2]="$NIXOS_USER"
  [nix-minimal]="$NIXOS_USER"
  [nixos]="$NIXOS_USER"
  [pxe-boot]="$NIXOS_USER"
  [server]="$NIXOS_USER"
  [tailscale-router]="$NIXOS_USER"
  [tor-relay]="$NIXOS_USER"
  [pve1]="${PROXMOX_SSH_USER:-wayne}"
)

# ── argument parsing ──────────────────────────────────────────────────────────

DRY_RUN=0
FILTER=()

for arg in "$@"; do
  case "$arg" in
    --dry-run) DRY_RUN=1 ;;
    --*)       echo "Unknown option: $arg" >&2; exit 1 ;;
    *)         FILTER+=("$arg") ;;
  esac
done

# Resolve the set of (hostname, user) pairs to process.
declare -A TARGET_USERS
if [[ ${#FILTER[@]} -gt 0 ]]; then
  for h in "${FILTER[@]}"; do
    if [[ -z "${HOSTS[$h]+_}" ]]; then
      echo "Unknown hostname: $h (not in the gc-hosts list)" >&2
      exit 1
    fi
    TARGET_USERS[$h]="${HOSTS[$h]}"
  done
else
  for h in "${!HOSTS[@]}"; do
    TARGET_USERS[$h]="${HOSTS[$h]}"
  done
fi

SORTED_HOSTS=($(printf '%s\n' "${!TARGET_USERS[@]}" | sort))

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

gc_one() {
  local host="$1" user="$2" logfile="$3"
  local target="${user}@${host}"

  if ! ssh "${SSH_OPTS[@]}" "$target" "true" 2>>"$logfile"; then
    echo "unreachable"
    return
  fi

  # NixOS hosts: try passwordless sudo first.
  if [[ "$user" == "$NIXOS_USER" ]]; then
    if ssh "${SSH_OPTS[@]}" "$target" "sudo -n nix-collect-garbage -d" \
        >>"$logfile" 2>>"$logfile"; then
      echo "ok(sudo)"
      return
    fi
    # sudo required a password — fall back to user-level gc.
    echo "[sudo needs password, falling back to user-level gc]" >>"$logfile"
    if ssh "${SSH_OPTS[@]}" "$target" "nix-collect-garbage -d" \
        >>"$logfile" 2>>"$logfile"; then
      echo "ok(user)"
      return
    fi
  else
    # Non-NixOS node (pve1): no system generations; just gc as the login user.
    if ssh "${SSH_OPTS[@]}" "$target" "nix-collect-garbage -d" \
        >>"$logfile" 2>>"$logfile"; then
      echo "ok"
      return
    fi
  fi

  echo "failed:$?"
}

# ── dry-run ───────────────────────────────────────────────────────────────────

if [[ "$DRY_RUN" -eq 1 ]]; then
  echo "[dry-run] would run gc on: ${SORTED_HOSTS[*]}"
  for host in "${SORTED_HOSTS[@]}"; do
    user="${TARGET_USERS[$host]}"
    if [[ "$user" == "$NIXOS_USER" ]]; then
      echo "  ssh ${SSH_OPTS[*]} ${user}@${host} 'sudo -n nix-collect-garbage -d'"
      echo "  # fallback: ssh ... 'nix-collect-garbage -d'"
    else
      echo "  ssh ${SSH_OPTS[*]} ${user}@${host} 'nix-collect-garbage -d'"
    fi
  done
  exit 0
fi

# ── parallel execution ────────────────────────────────────────────────────────

TMPDIR_GC="$(mktemp -d)"
trap 'rm -rf "$TMPDIR_GC"' EXIT

declare -A LOGS

echo "Running gc on ${#SORTED_HOSTS[@]} hosts (up to ${MAX_JOBS} parallel)..."
echo ""

job_count=0
for host in "${SORTED_HOSTS[@]}"; do
  user="${TARGET_USERS[$host]}"
  logfile="${TMPDIR_GC}/${host}.log"
  resultfile="${TMPDIR_GC}/${host}.result"
  LOGS[$host]="$logfile"
  : > "$logfile"

  (
    result="$(gc_one "$host" "$user" "$logfile")"
    echo "$result" > "$resultfile"
  ) &

  (( job_count++ )) || true
  if [[ "$job_count" -ge "$MAX_JOBS" ]]; then
    wait -n 2>/dev/null || wait
    (( job_count-- )) || true
  fi
done

wait

# ── summary ───────────────────────────────────────────────────────────────────

echo "Results:"
echo "──────────────────────────────"

ok_hosts=()
warn_hosts=()
fail_hosts=()

for host in "${SORTED_HOSTS[@]}"; do
  resultfile="${TMPDIR_GC}/${host}.result"
  result="$(cat "$resultfile" 2>/dev/null || echo "failed:missing")"

  case "$result" in
    ok|"ok(sudo)"|"ok(user)")
      printf "  %-20s  %s\n" "$host" "$result"
      ok_hosts+=("$host")
      ;;
    unreachable)
      printf "  %-20s  UNREACHABLE\n" "$host"
      warn_hosts+=("$host")
      ;;
    *)
      printf "  %-20s  FAILED (%s)\n" "$host" "$result"
      fail_hosts+=("$host")
      ;;
  esac
done

echo ""
echo "  ${#ok_hosts[@]} succeeded, ${#warn_hosts[@]} unreachable, ${#fail_hosts[@]} failed"

# Print logs for any non-ok host.
for host in "${warn_hosts[@]+"${warn_hosts[@]}"}" "${fail_hosts[@]+"${fail_hosts[@]}"}"; do
  logfile="${LOGS[$host]}"
  if [[ -s "$logfile" ]]; then
    echo ""
    echo "── $host ──"
    cat "$logfile"
  fi
done

echo ""
[[ "${#fail_hosts[@]}" -eq 0 ]]