#!/usr/bin/env bash
# One-command training round once labels exist (classifier-plan §3, §4.4, §5). Run on this host, with someone present.
#   training/runpod/runbook.sh BUNDLE_DIR --price 2.69 --image TAG [--volume ID] [--cloud COMMUNITY|SECURE] [--launch]
# Without --launch it runs the local checks, builds the robustness inputs and prints the plan; nothing is rented.
# With --launch:
#   pod.py create --yes -> ledger start -> wait for SSH -> sync up (+ watchdog cron) -> for each stage: ledger check,
#   run on the pod under tmux, sync down: setup -> smoke -> (ledger replan) -> screen -> final ARM -> evaluate ->
#   cleanup -> sync down --final (CLEAN verified) -> pod terminate -> local evaluation + acceptance.
# Ledger verdicts: 1 (projection past $35) = no new stage: the runbook stops and leaves the pod for the person present;
# 2 ($45 reached) = sync down, cleanup, terminate. Needs RUNPOD_API_KEY in the environment (never printed) and the
# RunPod SSH key (SSH_KEY, default ~/.ssh/naver_psy_runpod). Nothing is pushed to HF here (export_hf.py is a separate,
# explicit step after the acceptance review).
set -euo pipefail
HERE=$(cd "$(dirname "$0")" && pwd)
REPO=$(cd "$HERE/../.." && pwd)
PY=${PY:-$REPO/.venv/bin/python}
bundle=${1:-}; shift || true
price="" image="" volume="" cloud=COMMUNITY launch=0
while [ $# -gt 0 ]; do
  case "$1" in
    --price) price=$2; shift 2 ;;
    --image) image=$2; shift 2 ;;
    --volume) volume=$2; shift 2 ;;
    --cloud) cloud=$2; shift 2 ;;
    --launch) launch=1; shift ;;
    *) sed -n 2,13p "$0"; exit 2 ;;
  esac
done
[ -n "$bundle" ] && [ -f "$bundle/manifest.json" ] && [ -n "$price" ] && [ -n "$image" ] || { sed -n 2,13p "$0"; exit 2; }
bundle=$(cd "$bundle" && pwd)
name=$(basename "$bundle")
LOCAL=$HOME/naver_psy_models/$name
LEDGER=$LOCAL/ledger.json
mkdir -p "$LOCAL/inputs" && chmod 700 "$LOCAL"
ledger() { "$PY" -m training.ledger "$@" --ledger "$LEDGER"; }
pod() { "$PY" "$HERE/pod.py" "$@"; }

echo "== local checks"
(cd "$REPO" && "$PY" -m pytest -q tests/test_training_*.py tests/test_classifier_*.py)
[ -f "$LOCAL/inputs/robustness/robustness.jsonl" ] || \
  (cd "$REPO" && "$PY" -m training.robustness build --data "$bundle" --out "$LOCAL/inputs/robustness")
chmod -R go-rwx "$LOCAL/inputs"

cat <<EOF
== plan
bundle      $bundle
pod         1x H100 80GB, image $image, cloud $cloud, volume ${volume:-new 100 GB}, price \$$price/h
stages      setup 1.0h, smoke 0.4h, screen 0.8h (4 arms), final ~0.5h, evaluate 0.8h, cleanup  (≈ 3-3.5 h, \$9-14)
guards      ledger $LEDGER: no new stage past \$35, stop at \$45, cap \$50
results     $LOCAL/runs/
EOF
if [ $launch = 0 ]; then echo "dry run: add --launch to rent the pod"; exit 0; fi
[ ! -e "$LEDGER" ] || { echo "ledger exists: $LEDGER (one ledger per round)"; exit 1; }

echo "== create pod"
created=$(pod create --image "$image" --cloud "$cloud" ${volume:+--volume "$volume"} --yes)
POD=$("$PY" -c 'import json,sys; print(json.loads(sys.argv[1])["id"])' "$created")
started=$(date -u +%Y-%m-%dT%H:%M:%SZ)
ledger start --price "$price" --started-at "$started"
trap 'echo "!! runbook stopped; pod $POD may still be running: training/runpod/pod.py status $POD"' ERR
echo "pod $POD created at $started"

addr=""
for _ in $(seq 1 80); do
  if addr=$(pod ssh "$POD"); then break; fi
  sleep 15
done
[ -n "$addr" ] || { echo "no SSH address after 20 min"; exit 1; }
read -r HOST PORT <<<"$addr"
KEY=${SSH_KEY:-$HOME/.ssh/naver_psy_runpod}
SSH="ssh -p $PORT -i $KEY -o IdentitiesOnly=yes -o StrictHostKeyChecking=accept-new root@$HOST"
for _ in $(seq 1 40); do $SSH true 2>/dev/null && break; sleep 15; done

ROBUSTNESS=$LOCAL/inputs/robustness "$HERE/sync.sh" up "$HOST" "$PORT" "$bundle"
$SSH "(crontab -l 2>/dev/null; echo '*/5 * * * * RUNPOD_POD_ID=$POD /workspace/code/training/runpod/watchdog.sh >> /workspace/watchdog.log 2>&1') | crontab -" \
  || echo "WARN: watchdog cron not installed"

finish() {  # sync down, cleanup, verify CLEAN, terminate
  "$HERE/sync.sh" down "$HOST" "$PORT" "$name" || true
  remote cleanup
  "$HERE/sync.sh" down "$HOST" "$PORT" "$name" --final
  pod terminate "$POD" --yes
  ledger record --stage terminate --event note --note "pod $POD terminated"
}

remote() {  # run.sh STAGE [ARG] on the pod under tmux; waits for its exit code
  local tag; tag=$(printf '%s' "$1${2:+-$(basename "$2")}" | tr -c 'A-Za-z0-9-' '_')
  $SSH "cd /workspace/code && rm -f /workspace/$tag.rc && tmux new-session -d -s '$tag' \
    'EPOCHS=${EPOCHS:-2} bash training/runpod/run.sh $* > /workspace/$tag.log 2>&1; echo \$? > /workspace/$tag.rc'"
  until $SSH "test -f /workspace/$tag.rc" 2>/dev/null; do sleep 60; done
  local rc; rc=$($SSH "cat /workspace/$tag.rc")
  [ "$rc" = 0 ] || { echo "stage $tag failed (rc $rc); log: /workspace/$tag.log"; return 1; }
}

stage() {  # stage NAME HOURS [ARG]
  local rc=0
  ledger check --stage "$1" --hours "$2" || rc=$?
  if [ $rc = 2 ]; then echo "ledger: \$45 stop line"; finish; exit 1; fi
  if [ $rc = 1 ]; then echo "ledger: no new stage past \$35; pod $POD left running for a decision"; exit 1; fi
  ledger record --stage "$1" --event start
  remote "$1" ${3:+"$3"}
  ledger record --stage "$1" --event end
  "$HERE/sync.sh" down "$HOST" "$PORT" "$name"
}

stage setup 1.0
stage smoke 0.4
tps=$("$PY" - "$LOCAL/runs/smoke/train_log.jsonl" <<'EOF'
import json, sys
rows = [json.loads(l) for l in open(sys.argv[1]) if l.strip()]
rates = [r["tokens_per_second"] for r in rows if r.get("event") == "train" and r.get("tokens_per_second")]
print(sorted(rates)[len(rates) // 2] if rates else 0)
EOF
)
tokens=$("$PY" - "$bundle/train.jsonl" <<'EOF'
import sys
print(sum(1 for line in open(sys.argv[1]) if line.strip()) * 200)  # §4.5 assumption: ~200 tokens/post
EOF
)
arm=B-last epochs=2 rc=0
ledger replan --tokens-per-second "$tps" --train-tokens "$tokens" || rc=$?
if [ $rc = 0 ]; then
  stage screen 0.8
  arm=$("$PY" -c 'import json,sys; print(json.load(open(sys.argv[1]))["winner"])' "$LOCAL/runs/screen.json")
else
  epochs=1.5
  echo "replan: screen dropped, B-last, $epochs epochs"
fi
echo "arm $arm"
EPOCHS=$epochs stage final 1.0 "$arm"
stage evaluate 0.8 "/workspace/runs/final-$arm"
finish
trap - ERR

echo "== local evaluation"
RUN=$LOCAL/runs/final-$arm
cd "$REPO"
"$PY" -m training.evaluate --data "$bundle" --run "$RUN" --out "$RUN/metrics.json"
"$PY" -m training.robustness score --rows "$LOCAL/inputs/robustness/robustness.jsonl" \
  --base "$RUN/predictions/test.jsonl" --attacked "$RUN/predictions/robustness.jsonl" \
  --taxonomy "$RUN/taxonomy.json" --temperatures "$RUN/temperatures.json" --out "$RUN/robustness_score.json"
# Serving parity (§6.4) is measured later against node12 (training.parity); until then item 6 stays "missing",
# so this first decision can be no-go only for that reason. Rerun acceptance with --parity once it exists.
rc=0
"$PY" -m training.acceptance --metrics "$RUN/metrics.json" --robustness "$RUN/robustness_score.json" \
  --determinism "$RUN/test_predict.json" --out "$RUN/REPORT.md" || rc=$?
echo "done: $RUN (acceptance exit $rc: 0 go, 1 no-go, 2 user decision; REPORT.md + REPORT.json)"
echo "HF push is a separate, explicit step: python -m training.export_hf $RUN --out DIR --push --yes"