#!/usr/bin/env bash
# run_probe.sh — short-lived fit probe for the dual-gpu-tuner skill.
#
# Launches a COPY of the production launcher (never edits it) with a probe
# port and an optional -ot override, waits for the listening state, snapshots
# rocm-smi, kills the probe, parses the log, and appends a normalized record
# to the placement history.
#
# Usage:
#   run_probe.sh LAUNCHER [OT_OVERRIDE ...] [options]
#
#   OT_OVERRIDE   one or more, e.g. '^blk\.36\.\w[\w.]*$=ROCm1'   (optional)
#   --port N          probe port (default 8123)
#   --timeout N       seconds to wait for listening state (default 600)
#   --fit-target N    record a fit target different from the launcher's
#   --history PATH    history file (default <skill>/data/placement_history.jsonl)
#   --force           proceed even if a server for this model is already running
#   --keep            leave the probe server running after capture
#   --dry-run         print the probe script and exit
#
# The launcher's own server must be stopped first (the probe needs the GPUs).
# Launcher args must not contain a literal '|' character.

set -euo pipefail

SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
SKILL_DIR=$(dirname "$SCRIPT_DIR")
HISTORY="$SKILL_DIR/data/placement_history.jsonl"
PORT=8123
TIMEOUT=600
FIT_TARGET_OVERRIDE=""
OTS=()
FORCE=0
KEEP=0
DRY_RUN=0

usage() { sed -n '2,20p' "$0" | sed 's/^# \{0,1\}//'; exit 2; }

LAUNCHER=""
while [ $# -gt 0 ]; do
  case "$1" in
    --port) PORT="$2"; shift 2 ;;
    --timeout) TIMEOUT="$2"; shift 2 ;;
    --fit-target) FIT_TARGET_OVERRIDE="$2"; shift 2 ;;
    --history) HISTORY="$2"; shift 2 ;;
    --force) FORCE=1; shift ;;
    --keep) KEEP=1; shift ;;
    --dry-run) DRY_RUN=1; shift ;;
    -h|--help) usage ;;
    -*) echo "unknown option: $1" >&2; usage ;;
    *) if [ -z "$LAUNCHER" ]; then LAUNCHER="$1"; else OTS+=("$1"); fi; shift ;;
  esac
done
OT_STR=$(IFS=' + '; printf '%s' "${OTS[*]:-}")
[ -n "$LAUNCHER" ] || usage
[ -f "$LAUNCHER" ] || { echo "error: launcher not found: $LAUNCHER" >&2; exit 2; }

# ---------------------------------------------------------------- extraction
raw=$(grep -vE '^[[:space:]]*(#|$)' "$LAUNCHER" | sed -E 's/\\[[:space:]]*$//')
envs=$(grep -E '^[[:space:]]*export ' "$LAUNCHER" || true)
# strip the launcher's own -ot/--tensor-overrides BEFORE the '|' extraction:
# override values may contain '|' in regex alternations, which would truncate
# the extracted command (value may be single-quoted, double-quoted, or bare)
cmdline=$(printf '%s\n' "$raw" | grep -vE '^[[:space:]]*export ' | tr '\n' ' ' \
  | sed -E "s/(^|[[:space:]])(-ot|--tensor-overrides)[[:space:]]+('[^']*'|\"[^\"]*\"|[^[:space:]]+)//g")
cmd=$(printf '%s' "$cmdline" | grep -oE '[^|]*llama-server[^|]*' | head -1 || true)
# trim
cmd="${cmd#"${cmd%%[![:space:]]*}"}"
cmd="${cmd%"${cmd##*[![:space:]]}"}"
# strip trailing 2>&1 (monitor pipe)
cmd="${cmd% 2>&1}"
cmd="${cmd%"${cmd##*[![:space:]]}"}"
[ -n "$cmd" ] || { echo "error: could not extract llama-server command from $LAUNCHER" >&2; exit 2; }

MODEL=$(printf '%s' "$cmd" | grep -oE '\-m [^ ]+' | head -1 | awk '{print $2}')
MODEL_BASE=$(basename "${MODEL:-unknown}")

# safety: a server for this model must not be running (GPUs are busy)
if [ "$FORCE" -ne 1 ] && pgrep -f "llama-server.*$MODEL_BASE" >/dev/null 2>&1; then
  echo "error: a llama-server for $MODEL_BASE is already running." >&2
  echo "       stop it first (the probe needs the GPUs), or pass --force." >&2
  exit 3
fi
if ss -ltn 2>/dev/null | grep -qE "[:.]$PORT[[:space:]]"; then
  echo "error: port $PORT is already in use; pass --port N." >&2
  exit 3
fi

# ------------------------------------------------------------------ key json
KEY=$(printf '%s' "$cmd" | python3 -c '
import re, sys, json
s = sys.stdin.read()
def g(p):
    m = re.search(p, s)
    return m.group(1) if m else None
key = {
    "model_file": (g(r"-m (\S+)") or "").rsplit("/", 1)[-1],
    "tensor_split": g(r"--tensor-split (\S+)"),
    "ctk": g(r"-ctk (\S+)"), "ctv": g(r"-ctv (\S+)"),
    "ctk_draft": g(r"--cache-type-k-draft (\S+)"),
    "ctv_draft": g(r"--cache-type-v-draft (\S+)"),
    "vec": g(r"--hip-fa-force-vec (on|off)"),
    "fit_target": g(r"--fit-target (\S+)"),
    "devices": g(r"--device (\S+)"),
    "batch": (g(r"-b (\S+)") or "") + "/" + (g(r"-ub (\S+)") or ""),
    "ctx_checkpoints": g(r"--ctx-checkpoints (\S+)"),
}
print(json.dumps({k: v for k, v in key.items() if v}))
')
[ -n "$FIT_TARGET_OVERRIDE" ] && KEY=$(python3 -c '
import json,sys
k=json.loads(sys.argv[1]); k["fit_target"]=sys.argv[2]; print(json.dumps(k))' "$KEY" "$FIT_TARGET_OVERRIDE")

# -------------------------------------------------------------- probe script
TS=$(date +%Y%m%d_%H%M%S)
STAMP=$(date +%H%M%S)
WORKDIR=$(mktemp -d /tmp/otprobe.XXXXXX)
PROBE="$WORKDIR/probe_$STAMP.sh"
LOG="$WORKDIR/probe_$STAMP.log"
SMI="$WORKDIR/probe_$STAMP.rocm"
RUNJSON="$WORKDIR/probe_$STAMP.json"

{
  printf '#!/usr/bin/env bash\n'
  [ -n "$envs" ] && printf '%s\n' "$envs"
  printf 'exec %s' "$cmd"   # exec: wrapper PID becomes the server PID (kill works)
  printf ' --port %d' "$PORT"
  for ot in ${OTS[@]+"${OTS[@]}"}; do printf ' -ot %q' "$ot"; done
  printf '\n'
} > "$PROBE"
chmod +x "$PROBE"

if [ "$DRY_RUN" -eq 1 ]; then
  echo "# probe script (would run, log -> $LOG):"
  cat "$PROBE"
  echo "# key: $KEY"
  [ -n "$OT_STR" ] && echo "# ot:  $OT_STR"
  exit 0
fi

echo "[probe] model=$MODEL_BASE port=$PORT ot=${OT_STR:-none}"
echo "[probe] log=$LOG"
bash "$PROBE" > "$LOG" 2>&1 &
PID=$!

LISTENING=0
SECONDS=0
while [ "$SECONDS" -lt "$TIMEOUT" ]; do
  if grep -q "listening on" "$LOG" 2>/dev/null; then LISTENING=1; break; fi
  if grep -qE "allocation failed|out of memory|GGML_ASSERT|segmentation fault" "$LOG" 2>/dev/null; then break; fi
  kill -0 "$PID" 2>/dev/null || break
  sleep 2
done

# rocm-smi snapshot while the probe is alive (authoritative free memory)
# (ROCm only — skipped for Vulkan/other backends; fitter budgets are used)
if [ "$LISTENING" -eq 1 ] && command -v rocm-smi >/dev/null 2>&1; then
  sleep 2
  rocm-smi --showmeminfo vram > "$SMI" 2>/dev/null || true
fi
[ -s "$SMI" ] || SMI=""

if [ "$KEEP" -ne 1 ]; then
  kill -TERM "$PID" 2>/dev/null || true
  pkill -TERM -P "$PID" 2>/dev/null || true
  for _ in $(seq 1 15); do kill -0 "$PID" 2>/dev/null || break; sleep 1; done
  kill -9 "$PID" 2>/dev/null || true
  pkill -9 -P "$PID" 2>/dev/null || true
  # safety sweep: nothing for this model may survive the probe
  pkill -9 -f "llama-server.*$MODEL_BASE" 2>/dev/null || true
  sleep 1
fi

# ------------------------------------------------------------- parse + record
python3 "$SCRIPT_DIR/parse_llama_log.py" "$LOG" \
  ${SMI:+--rocm-smi "$SMI"} \
  --key-json "$KEY" \
  --ot "$OT_STR" \
  --tensor-split "$(printf '%s' "$cmd" | grep -oE -- '--tensor-split [^ ]+' | awk '{print $2}')" \
  --append "$HISTORY" \
  --quiet > "$RUNJSON" 2>&1 || true

RUN=$(python3 - "$RUNJSON" <<'EOF'
import json, sys
try:
    d = json.load(open(sys.argv[1]))["run"]
except Exception:
    print("parse-failed"); sys.exit(0)
f = d.get("free_mib", {})
m = d.get("model_mib", {})
devs = [x for x in (d.get("devices") or m or {}) if x != "CPU"] or ["ROCm0", "ROCm1"]
print(f"ctx={d.get('ctx_fitted')} "
      f"free0={f.get(devs[0])} free1={f.get(devs[1] if len(devs) > 1 else '')} "
      f"model0={m.get(devs[0])} model1={m.get(devs[1] if len(devs) > 1 else '')} "
      f"load={'ok' if d.get('load_ok') else 'FAIL'} splits={d.get('graph_splits')}")
EOF
)

echo "[probe] $RUN"
echo "[probe] log:     $LOG"
[ -n "$SMI" ] && echo "[probe] rocm-smi: $SMI"
echo "[probe] record:  $HISTORY"
if [ "$LISTENING" -ne 1 ]; then
  echo "[probe] WARNING: probe did not reach listening state; last log lines:" >&2
  tail -5 "$LOG" >&2 || true
  exit 1
fi
exit 0