#!/usr/bin/env bash
# =============================================================================
# ONE COMMAND: 131k extension -> NIAH eval -> upload verified -> teardown verdict.
# Run this when the GPUs are free. Everything before it is already staged.
#
# Refuses to start if another process holds GPU memory -- this box is shared.
# =============================================================================
set -uo pipefail
cd "$(dirname "$0")"
export PATH=/root/Adhi/BluTrain/envv/bin:$PATH
L=go_logs; mkdir -p "$L"; M="$L/master.log"
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }

USED=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | awk '{s+=$1} END{print s}')
if [ "${USED:-0}" -gt 4000 ]; then
  say "REFUSING TO START: ${USED} MiB of GPU memory is in use by someone else."
  say "  per-GPU: $(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader | tr '\n' ' ')"
  say "  re-run this script once they finish."
  exit 1
fi
say "GPUs free (${USED} MiB in use) -- starting"

say "=== STAGE 1/3: 131k training ==="
./t128_train.sh || { say "training did not complete -- see t128_logs/master.log"; exit 1; }

say "=== STAGE 2/3: NIAH eval, 3 models x 6 lengths, all 8 GPUs ==="
RN=$(cat t128_logs/runnum 2>/dev/null || echo "")
FIN=""
[ -n "$RN" ] && FIN=$(ls -t checkpoints_bluscriptcp/rank_0/blumodelcp_run${RN}_step_*.ckpt 2>/dev/null | head -1)
[ -z "$FIN" ] && FIN=$(ls -t checkpoints_bluscriptcp/rank_0/blumodelcp_run*_step_*.ckpt 2>/dev/null | head -1)
[ -z "$FIN" ] && { say "no 131k checkpoint -- cannot evaluate"; exit 1; }
say "t128 = $(basename "$FIN")  $(stat -c%s "$FIN") B"
EXTRA="t128=$(basename "$FIN") t64=t64.ckpt base680=varybatch.ckpt" \
MODELS="t128 t64 base680" LENS="4096 8192 16384 32744 65512 131000" N=100 \
RULER_TASKS="niah_single" GPUS=0,1,2,3,4,5,6,7 \
  ./run_ruler.sh || say "eval reported problems -- see ruler_logs/master.log"

say "=== STAGE 3/3: upload + verify ==="
up(){ [ -f "$1" ] || { say "  MISSING $1"; return 1; }
  hf upload unparallelled/BluScriptCP "$1" "$2" >>"$L/upload.log" 2>&1
  python - "$1" "$2" <<'V'
import sys,os
from huggingface_hub import HfApi
w=os.path.getsize(sys.argv[1])
g={f.rfilename:(f.size or 0) for f in HfApi().repo_info("unparallelled/BluScriptCP",files_metadata=True).siblings}.get(sys.argv[2])
ok=(g==w); print(("  VERIFIED  " if ok else "  *** FAILED ")+f"{sys.argv[2]} local={w:,} remote={g}")
sys.exit(0 if ok else 1)
V
}
D=$(date -u +%Y-%m-%d); BAD=0
up "$FIN" "checkpoints/t128_680m_$(basename "$FIN" .ckpt).ckpt" || BAD=$((BAD+1))
up Data_Loader/Data/mix128/MANIFEST.txt data/mix128_MANIFEST.txt || BAD=$((BAD+1))
for f in t128_train.sh GO.sh dl128.sh run_ruler.sh; do up "$f" "scripts/$f" || BAD=$((BAD+1)); done
for g in t128_logs/master.log t128_logs/train.log ruler_logs/master.log; do
  [ -f "$g" ] && { cp -f "$g" "${g}.snap"; up "${g}.snap" "logs/128k_$(basename $g)" || BAD=$((BAD+1)); }; done
for d in ruler_len*; do [ -f "$d/summary.csv" ] || continue
  up "$d/summary.csv" "evals/$D/${d}_summary.csv" || BAD=$((BAD+1))
  [ -f "$d/depth_summary.csv" ] && { up "$d/depth_summary.csv" "evals/$D/${d}_depth.csv" || BAD=$((BAD+1)); }; done
say "=== TEARDOWN VERDICT ==="
[ "$BAD" -eq 0 ] && say "ALL VERIFIED -- SAFE TO TEAR DOWN" || say "*** $BAD failure(s) -- DO NOT TEAR DOWN ***"