#!/usr/bin/env bash
# refresh-server-info.sh — pull a fresh system-details.txt from servers.
#
# Hosts are discovered by listing `servers/*/` directory names. Each name
# is used as the SSH target, so put matching entries in ~/.ssh/config to
# customize user / port / identity.
#
# Fallback: if the dir name doesn't resolve (and ssh_config doesn't rewrite
# it), the script looks for `servers/<host>/ssh-target` (one line, containing
# an IP or `user@ip`) and uses that instead. This keeps the tool working
# from a fresh clone without requiring DNS or ssh_config setup.
#
# For each host:
#   1. Run scripts/server_inspect.sh on the remote via `ssh <target> 'bash -s'`.
#   2. Write output atomically to `servers/<host>/system-details.txt`.
#      A failed SSH/run never clobbers the previous good snapshot.
#
# Exit status is non-zero if any host failed.
#
# Usage:
#   scripts/refresh-server-info.sh                        show this help
#   scripts/refresh-server-info.sh all                    refresh every host
#   scripts/refresh-server-info.sh ana-docker             refresh one host
#   scripts/refresh-server-info.sh ana-docker ana-ml2     refresh several
#   scripts/refresh-server-info.sh --dry-run all          preview, no ssh
#   scripts/refresh-server-info.sh --validate-only all    checks only
#   scripts/refresh-server-info.sh --validate-only <h>    checks one host
#
# Discovery is validated before any ssh is attempted. Each host prints
# an indented "! <reason>" line per warning (unreadable files, empty or
# missing ssh-target, unresolvable name with no fallback, missing README
# or system-details, etc). Warnings never block the refresh — they're
# informational — but are surfaced so drift is visible.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
INSPECT="$SCRIPT_DIR/server_inspect.sh"
SERVERS_DIR="$REPO_ROOT/servers"

if [ ! -f "$INSPECT" ]; then
  echo "error: $INSPECT not found" >&2
  exit 2
fi

DRY_RUN=0
VALIDATE_ONLY=0
REQUESTED=()
for arg in "$@"; do
  case "$arg" in
    --dry-run) DRY_RUN=1 ;;
    --validate-only|--validate) VALIDATE_ONLY=1 ;;
    -h|--help)
      sed -n '2,32p' "$0"
      exit 0
      ;;
    -*) echo "error: unknown flag $arg" >&2; exit 2 ;;
    *)  REQUESTED+=("$arg") ;;
  esac
done

discover_hosts() {
  find "$SERVERS_DIR" -mindepth 1 -maxdepth 1 -type d -printf '%f\n' | sort
}

# No positional args → show help. Fleet-wide operations must be explicit.
if [ "${#REQUESTED[@]}" -eq 0 ]; then
  sed -n '2,32p' "$0"
  exit 0
fi

# The literal keyword "all" expands to every discovered host. Using it as
# a sentinel (rather than making "no args" mean "all") makes fleet-wide
# runs deliberate — useful when you might otherwise hit 20 hosts by accident.
if [ "${#REQUESTED[@]}" -eq 1 ] && [ "${REQUESTED[0]}" = "all" ]; then
  mapfile -t HOSTS < <(discover_hosts)
else
  # Check for 'all' mixed with other names — almost certainly a mistake.
  for r in "${REQUESTED[@]}"; do
    if [ "$r" = "all" ]; then
      echo "error: 'all' must be the only argument when used" >&2
      exit 2
    fi
  done
  HOSTS=("${REQUESTED[@]}")
fi

if [ "${#HOSTS[@]}" -eq 0 ]; then
  echo "error: no hosts found under $SERVERS_DIR" >&2
  exit 2
fi

resolve_target() {
  # Prints the SSH target to use for a given dir name.
  # Preference order:
  #   1. servers/<host>/ssh-target if present — it may carry a user
  #      prefix (root@host) or port that /etc/hosts + ssh_config can't
  #      express, so an explicit file is authoritative when it exists.
  #   2. The dir name, if /etc/hosts or DNS resolves it.
  #   3. The dir name as-is — let ssh fail with its own error.
  local host="$1"
  local fallback="$SERVERS_DIR/$host/ssh-target"
  if [ -f "$fallback" ]; then
    local target
    target=$(awk 'NF{print $1; exit}' "$fallback")
    if [ -n "$target" ]; then echo "$target"; return; fi
  fi
  local effective
  effective=$(ssh -G "$host" 2>/dev/null | awk '/^hostname /{print $2; exit}')
  if [ -n "$effective" ] && getent hosts "$effective" >/dev/null 2>&1; then
    echo "$host"
    return
  fi
  echo "$host"
}

validate_host() {
  # Prints one warning per line. Empty output = clean.
  local host="$1"
  local dir="$SERVERS_DIR/$host"

  if [ ! -d "$dir" ]; then
    printf '%s\n' "dir missing: $dir"
    return
  fi
  if [ ! -r "$dir" ] || [ ! -x "$dir" ]; then
    printf '%s\n' "dir not readable/searchable (check permissions)"
    return
  fi

  local stf="$dir/ssh-target"
  if [ -e "$stf" ]; then
    if [ ! -f "$stf" ]; then
      printf '%s\n' "ssh-target is not a regular file"
    elif [ ! -r "$stf" ]; then
      printf '%s\n' "ssh-target not readable"
    elif [ ! -s "$stf" ]; then
      printf '%s\n' "ssh-target is empty"
    else
      local t
      t=$(awk 'NF{print $1; exit}' "$stf" 2>/dev/null || true)
      if [ -z "$t" ]; then
        printf '%s\n' "ssh-target has no non-blank content"
      elif [[ "$t" =~ [[:space:]] ]]; then
        printf '%s\n' "ssh-target first token contains whitespace ('$t')"
      fi
    fi
  fi

  local effective=""
  effective=$(ssh -G "$host" 2>/dev/null | awk '/^hostname /{print $2; exit}')
  local resolves=0
  if [ -n "$effective" ] && getent hosts "$effective" >/dev/null 2>&1; then
    resolves=1
  fi
  if [ "$resolves" -eq 0 ] && [ ! -s "$stf" ]; then
    printf '%s\n' "name '$host' does not resolve and no ssh-target fallback present"
  fi

  [ -f "$dir/README.md" ]         || printf '%s\n' "README.md missing"
  if [ ! -f "$dir/system-details.txt" ]; then
    printf '%s\n' "system-details.txt missing (never refreshed)"
  else
    validate_snapshot "$dir/system-details.txt"
  fi
}

validate_snapshot() {
  # Read the captured system-details.txt and yield a warning for each
  # known problem marker that indicates the remote inspect ran but didn't
  # have the rights / environment it needed.
  local file="$1"
  if grep -q 'docker daemon not reachable by current user' "$file"; then
    printf '%s\n' "remote user cannot reach docker daemon (add to 'docker' group or fix socket perms)"
  fi
  if grep -qE '^Server:[[:space:]]*$' "$file" \
     && ! grep -q 'docker daemon not reachable by current user' "$file"; then
    printf '%s\n' "docker 'Server:' line is blank (daemon down or inaccessible)"
  fi
  if grep -q '^nvidia-smi present, but failed' "$file"; then
    printf '%s\n' "nvidia-smi failed on remote (driver broken or no permission)"
  fi
  if ! grep -q '^===== DONE =====' "$file"; then
    printf '%s\n' "snapshot appears truncated (missing trailing '===== DONE =====' marker)"
  fi
}

print_warnings() {
  # $1 = indent, rest = warning lines
  local indent="$1"; shift
  local w
  for w in "$@"; do
    printf '%s! %s\n' "$indent" "$w"
  done
}

pad=0
for h in "${HOSTS[@]}"; do (( ${#h} > pad )) && pad=${#h}; done

if [ "$VALIDATE_ONLY" -eq 1 ]; then
  printf 'Validating %d host(s):\n' "${#HOSTS[@]}"
  total_warn=0
  for host in "${HOSTS[@]}"; do
    mapfile -t warnings < <(validate_host "$host")
    if [ "${#warnings[@]}" -eq 0 ]; then
      printf '  %-*s  ok\n' "$pad" "$host"
    else
      printf '  %-*s  %d warning(s)\n' "$pad" "$host" "${#warnings[@]}"
      print_warnings "      " "${warnings[@]}"
      total_warn=$((total_warn + ${#warnings[@]}))
    fi
  done
  if [ "$total_warn" -gt 0 ]; then exit 1; fi
  exit 0
fi

printf 'Refreshing %d host(s):\n' "${#HOSTS[@]}"
failed=()
for host in "${HOSTS[@]}"; do
  out="$SERVERS_DIR/$host/system-details.txt"
  tmp="$out.new"

  mapfile -t warnings < <(validate_host "$host")

  target=$(resolve_target "$host")
  if [ "$target" = "$host" ]; then
    label="$host"
  else
    label="$host → $target"
  fi

  printf '  %-*s  ' "$pad" "$label"

  if [ "$DRY_RUN" -eq 1 ]; then
    printf 'would run: ssh %s bash -s < %s > %s\n' "$target" "$INSPECT" "$out"
    [ "${#warnings[@]}" -gt 0 ] && print_warnings "      " "${warnings[@]}"
    continue
  fi

  mkdir -p "$SERVERS_DIR/$host"

  if ssh -o BatchMode=yes -o ConnectTimeout=10 -o StrictHostKeyChecking=accept-new "$target" 'bash -s' < "$INSPECT" > "$tmp" 2> "$tmp.err"; then
    mv "$tmp" "$out"
    rm -f "$tmp.err"
    bytes=$(wc -c < "$out")
    printf 'ok (%s bytes)\n' "$bytes"
  else
    rc=$?
    rm -f "$tmp"
    err=$(head -n 1 "$tmp.err" 2>/dev/null || true)
    rm -f "$tmp.err"
    printf 'FAIL (rc=%d) %s\n' "$rc" "$err"
    failed+=("$host")
  fi

  [ "${#warnings[@]}" -gt 0 ] && print_warnings "      " "${warnings[@]}"
done

if [ "${#failed[@]}" -gt 0 ]; then
  printf '\n%d host(s) failed: %s\n' "${#failed[@]}" "${failed[*]}" >&2
  exit 1
fi
