#!/usr/bin/env bash
# gc-hosts.sh — Run nix-collect-garbage -d on all live NixOS hosts.
#
# The host list is rebuilt on every run:
#   1. This workstation (nixos) — always first
#   2. pve1 — always second (non-NixOS Proxmox node with Nix installed)
#   3. Every NixOS guest currently running on pve1 (discovered via pct/qm list)
#
# nix-cache is excluded: gc-ing the shared binary cache evicts store paths
# that other hosts depend on for substitution.
#
# NixOS hosts: tries "sudo -n nix-collect-garbage -d" first (works when
# wheelNeedsPassword = false, e.g. the HA cluster).  Falls back to user-level
# "nix-collect-garbage -d" if sudo needs a password — still collects
# unreferenced store paths and old nixos-user profile generations, but leaves
# old system generations in place.
# pve1: runs "bash -l -c nix-collect-garbage -d" as the login user so
# /etc/profile is sourced and the Nix daemon's PATH is set up automatically.
#
# Usage (from repo root):
#   bash scripts/gc-hosts.sh [--dry-run]

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

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

: "${MAX_JOBS:=8}"
: "${NIXOS_USER:=nixos}"
: "${PVE1_SSH_USER:=${PROXMOX_SSH_USER:-wayne}}"

# GC connections use BatchMode — no interactive prompts, just succeed or fail.
SSH_OPTS=(-o StrictHostKeyChecking=no -o BatchMode=yes -o ConnectTimeout=10)
# Discovery connections do NOT use BatchMode so that sudo can prompt if needed
# (pct/qm list require root access on Proxmox).
SSH_QUERY_OPTS=(-o StrictHostKeyChecking=no -o ConnectTimeout=10)

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

# ── build the host list ───────────────────────────────────────────────────────

# ORDERED_HOSTS: names in display/execution order.
# HOST_TARGET[name]: SSH target string (user@host).
# HOST_TYPE[name]: "nixos" (try sudo gc, fallback user) | "nix" (login-shell gc).
declare -a  ORDERED_HOSTS=()
declare -A  HOST_TARGET=()
declare -A  HOST_TYPE=()
declare -A  _SEEN_HOSTNAMES=()   # dedup tracker

_add_host() {
  local name="$1" target="$2" type="$3"
  if [[ -n "${_SEEN_HOSTNAMES[$name]+_}" ]]; then return; fi
  _SEEN_HOSTNAMES[$name]=1
  ORDERED_HOSTS+=("$name")
  HOST_TARGET[$name]="$target"
  HOST_TYPE[$name]="$type"
}

# 1. Workstation (hard-wired first)
_add_host "nixos" "${NIXOS_USER}@nixos" "nixos"

# 2. pve1 (hard-wired second; non-NixOS, no system generations)
_add_host "pve1" "${PVE1_SSH_USER}@${PVE1_HOST}" "nix"

# 3. Dynamically discover running NixOS guests on pve1
echo "Discovering running guests on ${PVE1_HOST}..."

# Evaluate the full flake hostname map in one shot.
hostname_map="{}"
if ! hostname_map="$(
  nix eval --json "${NIX_EVAL_FLAGS[@]}" .#nixosConfigurations \
    --apply 'cfgs: builtins.mapAttrs (_: cfg: cfg.config.networking.hostName) cfgs' \
    2>/dev/null
)"; then
  echo "  warning: flake eval failed — skipping dynamic host discovery" >&2
fi

# SSH_QUERY_OPTS (no BatchMode) so sudo can prompt if wayne's sudo needs a password.
if ssh "${SSH_QUERY_OPTS[@]}" "${PVE1_SSH_USER}@${PVE1_HOST}" "true" 2>/dev/null; then
  running_guests="$(
    ssh "${SSH_QUERY_OPTS[@]}" "${PVE1_SSH_USER}@${PVE1_HOST}" bash <<'REMOTE'
      { sudo pct list 2>/dev/null | awk 'NR>1 && $2=="running" { print $NF }';
        sudo qm  list 2>/dev/null | awk 'NR>1 && $3=="running" { print $2  }'; } | sort -u
REMOTE
  )" || running_guests=""

  while IFS= read -r guest; do
    [[ -z "$guest" ]] && continue

    # Resolve flake target name → NixOS hostname.
    hostname="$(printf '%s' "$hostname_map" \
      | jq -r --arg g "$guest" '.[$g] // empty' 2>/dev/null || true)"
    [[ -z "$hostname" ]] && continue

    # Exclude nix-cache and any hostname already in the list.
    case "$hostname" in nix-cache) continue ;; esac
    if [[ -n "${_SEEN_HOSTNAMES[$hostname]+_}" ]]; then continue; fi

    echo "  + $guest → $hostname"
    _add_host "$hostname" "${NIXOS_USER}@${hostname}" "nixos"
  done <<< "$running_guests"
else
  echo "  warning: ${PVE1_HOST} unreachable — skipping dynamic host discovery" >&2
fi

echo ""
echo "Hosts: ${ORDERED_HOSTS[*]}"
echo ""

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

if [[ "$DRY_RUN" -eq 1 ]]; then
  echo "[dry-run] commands that would run:"
  for host in "${ORDERED_HOSTS[@]}"; do
    target="${HOST_TARGET[$host]}"
    type="${HOST_TYPE[$host]}"
    if [[ "$type" == "nixos" ]]; then
      echo "  ssh ${SSH_OPTS[*]} $target 'sudo -n nix-collect-garbage -d'"
      echo "  # fallback: ssh ... $target 'nix-collect-garbage -d'"
    else
      echo "  ssh ${SSH_OPTS[*]} $target 'bash -l -c nix-collect-garbage -d'"
    fi
  done
  exit 0
fi

# ── gc worker ─────────────────────────────────────────────────────────────────

gc_one() {
  local host="$1" target="${HOST_TARGET[$1]}" type="${HOST_TYPE[$1]}" logfile="$2"

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

  if [[ "$type" == "nixos" ]]; then
    if ssh "${SSH_OPTS[@]}" "$target" "sudo -n nix-collect-garbage -d" \
        >>"$logfile" 2>>"$logfile"; then
      echo "ok(sudo)"; return
    fi
    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: use a login shell so /etc/profile is sourced and the
    # Nix daemon's bin dir is on PATH (set up by /etc/profile.d/nix-daemon.sh
    # which the Nix installer adds to /etc/profile).
    if ssh "${SSH_OPTS[@]}" "$target" "bash -l -c 'nix-collect-garbage -d'" \
        >>"$logfile" 2>>"$logfile"; then
      echo "ok"; return
    fi
  fi

  echo "failed:$?"
}

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

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

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

declare -A LOGS=()
job_count=0

for host in "${ORDERED_HOSTS[@]}"; do
  logfile="${TMPDIR_GC}/${host}.log"
  resultfile="${TMPDIR_GC}/${host}.result"
  LOGS[$host]="$logfile"
  : > "$logfile"

  ( result="$(gc_one "$host" "$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 "${ORDERED_HOSTS[@]}"; do
  result="$(cat "${TMPDIR_GC}/${host}.result" 2>/dev/null || echo "failed:missing")"
  case "$result" in
    ok|"ok(sudo)"|"ok(user)")
      printf "  %-22s  %s\n" "$host" "$result"
      ok_hosts+=("$host") ;;
    unreachable)
      printf "  %-22s  UNREACHABLE\n" "$host"
      warn_hosts+=("$host") ;;
    *)
      printf "  %-22s  FAILED (%s)\n" "$host" "$result"
      fail_hosts+=("$host") ;;
  esac
done

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

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 ]]