#!/bin/bash
# write_handoff.sh — assemble HANDOFF.md from the current run state.
#
# Sections:
#   - Leader checkpoint (highest bench_pass row)
#   - Per-arm summary (status, bench, val_loss, criteria pass/fail, rationale)
#   - Open threads (queued / running arms not yet terminal)
#   - Paste-ready next-session prompt
#
# Triggered by orchestrator on: SOTA improvement / queue done / budget cap.
#
# Usage: tools/orch/write_handoff.sh [--reason <text>]

set -u
SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)
REPO_ROOT=$(cd "$SCRIPT_DIR/../.." && pwd)
cd "$REPO_ROOT"

REASON="manual"
while [ $# -gt 0 ]; do
  case $1 in
    --reason) REASON=$2; shift 2;;
    *) shift;;
  esac
done

OUT=HANDOFF.md
TMP=$(mktemp)

NOW=$(date +%Y-%m-%dT%H:%M:%S%z)
{
  echo "# HANDOFF.md"
  echo ""
  echo "Auto-generated by \`tools/orch/write_handoff.sh\`. Triggered by: **${REASON}**."
  echo ""
  echo "Generated: ${NOW}"
  echo ""
} > "$TMP"

# ── Collect arm data into per-arm dicts (associative array via tmpfiles)
LEADER_PASS=-1
LEADER_ARM=""
LEADER_VL="-"
LEADER_CKPT="-"

declare -a BENCHED_ARMS
declare -a OPEN_ARMS
declare -a FAILED_ARMS
declare -a KILLED_ARMS

for run_dir in runs/*/; do
  rc=$run_dir/run_card.meta
  br=$run_dir/bench_result.meta
  [ -f "$rc" ] || continue
  arm=$(grep "^arm_id=" "$rc" | head -1 | cut -d= -f2-)
  status=$(grep "^status=" "$rc" | head -1 | cut -d= -f2-)
  case "$status" in
    queued|running)  OPEN_ARMS+=("$arm");;
    failed)          FAILED_ARMS+=("$arm");;
    killed)          KILLED_ARMS+=("$arm");;
    completed|benched) BENCHED_ARMS+=("$arm")
      if [ -f "$br" ]; then
        bp=$(grep "^bench_pass=" "$br" | head -1 | cut -d= -f2-)
        if [[ "$bp" =~ ^[0-9]+$ ]] && [ "$bp" -gt "$LEADER_PASS" ]; then
          LEADER_PASS=$bp
          LEADER_ARM=$arm
          ckpt=$(grep "^ckpt_prefix=" "$rc" | head -1 | cut -d= -f2-)
          LEADER_CKPT="${ckpt}_best"
          if [ -f "${ckpt}_best.meta" ]; then
            LEADER_VL=$(grep "^best_val_loss=" "${ckpt}_best.meta" | head -1 | cut -d= -f2-)
            LEADER_VL=$(printf '%.3f' "$LEADER_VL" 2>/dev/null || echo "$LEADER_VL")
          fi
        fi
      fi
      ;;
  esac
done

# ── Leader section ────────────────────────────────────────────────
{
  echo "## Leader"
  echo ""
  if [ "$LEADER_PASS" -ge 0 ]; then
    echo "- **Arm**: \`$LEADER_ARM\`"
    echo "- **Bench (single-best, strip-graded)**: $LEADER_PASS/30"
    echo "- **Best val_loss**: $LEADER_VL"
    echo "- **Checkpoint**: \`$LEADER_CKPT\`"
  else
    echo "- No benched arms yet."
  fi
  echo ""
} >> "$TMP"

# ── Per-arm summary ──────────────────────────────────────────────
{
  echo "## Per-arm summary"
  echo ""
  echo "| Arm | Status | Bench | Best val_loss | Notes |"
  echo "|-----|:------:|:-----:|--------------:|-------|"
  for run_dir in runs/*/; do
    rc=$run_dir/run_card.meta
    br=$run_dir/bench_result.meta
    [ -f "$rc" ] || continue
    arm=$(grep "^arm_id=" "$rc" | head -1 | cut -d= -f2-)
    status=$(grep "^status=" "$rc" | head -1 | cut -d= -f2-)
    rationale=$(grep "^rationale=" "$rc" | head -1 | cut -d= -f2- | head -c 60)
    bp="-"
    if [ -f "$br" ]; then
      bp=$(grep "^bench_pass=" "$br" | head -1 | cut -d= -f2-)
    fi
    ckpt=$(grep "^ckpt_prefix=" "$rc" | head -1 | cut -d= -f2-)
    vl="-"
    if [ -n "${ckpt:-}" ] && [ -f "${ckpt}_best.meta" ]; then
      vl=$(grep "^best_val_loss=" "${ckpt}_best.meta" | head -1 | cut -d= -f2-)
      vl=$(printf '%.3f' "$vl" 2>/dev/null || echo "$vl")
    fi
    echo "| $arm | $status | $bp | $vl | ${rationale}... |"
  done
  echo ""
} >> "$TMP"

# ── Falsified hypotheses (terminal arms that did NOT improve SOTA)
{
  echo "## Falsified hypotheses"
  echo ""
  echo "Terminal arms whose bench did not exceed leader. Per arm \`rationale\` is"
  echo "a hypothesis statement; bench < leader = the hypothesis is falsified at this scale."
  echo ""
  any=0
  for arm in "${BENCHED_ARMS[@]:-}"; do
    [ -n "$arm" ] || continue
    [ "$arm" = "$LEADER_ARM" ] && continue
    rc=runs/$arm/run_card.meta
    rationale=$(grep "^rationale=" "$rc" | head -1 | cut -d= -f2-)
    echo "- \`$arm\`: $rationale"
    any=1
  done
  [ "$any" = "0" ] && echo "(none)"
  echo ""
} >> "$TMP"

# ── Killed / failed
{
  echo "## Non-completed arms"
  echo ""
  if [ ${#KILLED_ARMS[@]:-0} -gt 0 ]; then
    echo "**Killed by σ-laggard rule:**"
    for arm in "${KILLED_ARMS[@]}"; do echo "- \`$arm\`"; done
    echo ""
  fi
  if [ ${#FAILED_ARMS[@]:-0} -gt 0 ]; then
    echo "**Failed (trainer error / OOM / etc):**"
    for arm in "${FAILED_ARMS[@]}"; do echo "- \`$arm\` — see \`runs/$arm/train.log\`"; done
    echo ""
  fi
  if [ ${#KILLED_ARMS[@]:-0} -eq 0 ] && [ ${#FAILED_ARMS[@]:-0} -eq 0 ]; then
    echo "(none)"
    echo ""
  fi
} >> "$TMP"

# ── Open threads
{
  echo "## Open threads"
  echo ""
  if [ ${#OPEN_ARMS[@]:-0} -gt 0 ]; then
    echo "Arms not yet terminal at handoff time:"
    for arm in "${OPEN_ARMS[@]}"; do
      rc=runs/$arm/run_card.meta
      st=$(grep "^status=" "$rc" | head -1 | cut -d= -f2-)
      echo "- \`$arm\` (status=$st)"
    done
    echo ""
  else
    echo "(none — queue exhausted)"
    echo ""
  fi
} >> "$TMP"

# ── Paste-ready next-session prompt
{
  echo "## Paste-ready next-session prompt"
  echo ""
  echo '```'
  echo "Resuming Spur orchestrator session. Current state:"
  echo ""
  if [ "$LEADER_PASS" -ge 0 ]; then
    echo "- Leader: $LEADER_ARM @ $LEADER_PASS/30 (val=$LEADER_VL, ckpt=$LEADER_CKPT)"
  else
    echo "- No benched arms yet."
  fi
  echo "- ${#OPEN_ARMS[@]:-0} open arm(s); ${#BENCHED_ARMS[@]:-0} terminal benched; ${#KILLED_ARMS[@]:-0} killed; ${#FAILED_ARMS[@]:-0} failed."
  echo ""
  echo "Next steps:"
  echo "1. Review LEADERBOARD.md for full table."
  echo "2. Run \`tools/orch/condense.rail\` to compute per-prompt max-pass ensemble across all benched ckpts."
  echo "3. If queue exhausted: design next batch of arms in EXPERIMENT_PLAN.md (per blob/slice/fan/condense methodology)."
  echo "4. If open arms remain: \`tools/orch/poll.sh runs/<arm_id>\` for each, decide kill/keep, possibly relaunch via \`launch_arm.sh\`."
  echo '```'
  echo ""
} >> "$TMP"

mv "$TMP" "$OUT"
echo "wrote $OUT (leader=$LEADER_ARM @ $LEADER_PASS/30; reason=$REASON)"
exit 0
