#!/usr/bin/env bash
# Progress check for remote GPU builder.
# Run while a pipeline build is in progress to monitor GPU/CPU utilization.
#
# Usage:
#   tools/builder-progress-check.sh                 # one-shot
#   tools/builder-progress-check.sh --watch         # repeat every 15s
#   tools/builder-progress-check.sh --watch --interval 5

set -euo pipefail

REPO_ROOT="$(git rev-parse --show-toplevel 2>/dev/null || pwd)"
STATE_FILE="${REPO_ROOT}/.devenv/state/builder-host.json"

if [ ! -f "$STATE_FILE" ]; then
  echo "No builder state at $STATE_FILE" >&2
  exit 1
fi

HOST=$(python3 -c "import json; print(json.load(open('$STATE_FILE'))['server_host'])")
SSH_OPTS="-o BatchMode=yes -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o ConnectTimeout=10 -o LogLevel=ERROR"

interval=0
while [[ "${1:-}" == --* ]]; do
  case "$1" in
    --watch) interval=15; shift ;;
    --interval) interval="$2"; shift 2 ;;
    *) echo "Unknown option: $1" >&2; exit 1 ;;
  esac
done

check() {
  echo "=== $(date -Iseconds) === builder@[$HOST]"
  echo ""

  # GPU utilization
  echo "--- GPU ---"
  ssh $SSH_OPTS "root@$HOST" 'nvidia-smi --query-gpu=name,utilization.gpu,utilization.memory,memory.used,memory.total,temperature.gpu,power.draw --format=csv,noheader 2>/dev/null || echo "nvidia-smi not available"'
  echo ""

  # CPU / load
  echo "--- CPU ---"
  ssh $SSH_OPTS "root@$HOST" 'nproc && cat /proc/loadavg'
  echo ""

  # Memory
  echo "--- Memory ---"
  ssh $SSH_OPTS "root@$HOST" 'free -h | head -2'
  echo ""

  # Disk
  echo "--- Disk ---"
  ssh $SSH_OPTS "root@$HOST" 'df -h /nix 2>/dev/null || df -h /'
  echo ""

  # Active Nix builds
  echo "--- Active Nix builds ---"
  ssh $SSH_OPTS "root@$HOST" 'ps aux --sort=-%cpu | grep -E "nix-build|python|ffmpeg|whisper" | grep -v grep | head -10 || echo "(none)"'
  echo ""

  # Nix daemon queue
  echo "--- Nix daemon ---"
  ssh $SSH_OPTS "root@$HOST" 'systemctl is-active nix-daemon.service 2>/dev/null; ls /nix/var/tmp/nix-build/ 2>/dev/null | wc -l | xargs -I{} echo "Active build dirs: {}"'
  echo ""
}

check

if [ "$interval" -gt 0 ]; then
  while true; do
    sleep "$interval"
    check
  done
fi