#!/usr/bin/env bash
# Helper for the QuickPod GPU box. The SSH target comes ONLY from $QUICKPOD_SSH_SECRET
# (format: "ssh user@host -p PORT"); it is never printed. Transfers use tar-over-ssh
# (rsync 3.5 <-> 3.2 protocol mismatch was observed).
# Usage:
#   scripts/quickpod.sh status [OUT]        driver/trainer PIDs, GPU, last log lines, last progress rows
#   scripts/quickpod.sh pull [OUT] [DEST]   copy checkpoints (+progress.tsv) not yet present locally
#   scripts/quickpod.sh latest [OUT] [DEST] pull only the newest checkpoint
#   scripts/quickpod.sh ssh 'cmd'           run a remote command (cwd ~/atom-ai)
# OUT = remote checkpoint dir relative to ~/atom-ai (default checkpoints/fr_big_b1024)
# DEST default: /workspace/exports/atom-ai-batched/checkpoints/quickpod/<basename OUT>  (existing files never overwritten)
set -euo pipefail
: "${QUICKPOD_SSH_SECRET:?QUICKPOD_SSH_SECRET not set}"
RDIR=${RDIR:-'$HOME/atom-ai'}
qp() { set -- $QUICKPOD_SSH_SECRET "$@"; shift
  ssh -o StrictHostKeyChecking=accept-new -o BatchMode=yes -o ConnectTimeout=20 -o LogLevel=ERROR "$@" \
    2> >(sed -E 's/[0-9]{1,3}(\.[0-9]{1,3}){3}/<IP>/g' >&2); }
cmd=${1:-status}; OUT=${2:-checkpoints/fr_big_b1024}
DEST=${3:-/workspace/exports/atom-ai-batched/checkpoints/quickpod/$(basename "$OUT")}
case "$cmd" in
  status)
    qp "cd $RDIR; echo '== driver'; cat $OUT/driver.pid 2>/dev/null && pgrep -fa 'python.*train_batche[d]' | cut -c1-200 || true
        echo '== gpu'; nvidia-smi --query-gpu=name,utilization.gpu,memory.used,memory.total --format=csv,noheader
        echo '== disk'; df -h . | tail -1
        echo '== log'; tail -n 8 logs/batched/\$(ls -t logs/batched | head -1) 2>/dev/null | cut -c1-300
        echo '== progress'; tail -n 5 $OUT/progress.tsv 2>/dev/null | cut -f1-12
        echo '== ckpts'; ls $OUT/*.pt 2>/dev/null | wc -l; ls -t $OUT/*.pt 2>/dev/null | head -1" ;;
  pull|latest)
    mkdir -p "$DEST"
    if [ "$cmd" = latest ]; then files=$(qp "cd $RDIR/$OUT && ls -t *.pt | head -1")
    else files=$(qp "cd $RDIR/$OUT && ls *.pt"); fi
    todo=(); for f in $files; do [ -e "$DEST/$f" ] || todo+=("$f"); done
    [ ${#todo[@]} -gt 0 ] && qp "cd $RDIR/$OUT && tar cf - ${todo[*]}" | tar xf - -C "$DEST" --skip-old-files
    qp "cd $RDIR/$OUT && cat progress.tsv 2>/dev/null" > "$DEST/progress.tsv" || true
    echo "pulled ${#todo[@]} checkpoint(s) into $DEST"; ls -t "$DEST"/*.pt 2>/dev/null | head -3 ;;
  ssh) shift; qp "cd $RDIR && $*" ;;
  *) sed -n '2,12p' "$0"; exit 1 ;;
esac