#!/bin/bash
# Usage: eval_checkpoint.sh <checkpoint-path> <tag> [n_tb2] [n_swe] [temp]
# Full loop: kill GPU procs, serve checkpoint, wait healthy, run both suites, shutdown.
set -u
CKPT=$1
TAG=$2
N_TB2=${3:-89}
N_SWE=${4:-100}
TEMP=${5:-}
PORT=8500

cd /root/work/a/prime-rl
export PRIME_API_KEY="$(cat "$AGENTPTB_PRIME_KEY_FILE")"

# kill anything holding GPUs
pkill -f "PRIME-RL" 2>/dev/null || true
sleep 5
pkill -9 -f "PRIME-RL" 2>/dev/null || true
sleep 3

# serve
CUDA_VISIBLE_DEVICES=0 nohup uv run inference \
  --model.name "$CKPT" \
  --model.tool-call-parser qwen3_coder \
  --model.reasoning-parser qwen3 \
  --server.port $PORT --parallel.tp 1 --parallel.dp 1 \
  --vllm-extra '{"limit_mm_per_prompt": {"image": 0, "video": 0}}' \
  > $AGENTPTB_WORKSPACE/logs/inference-$TAG.log 2>&1 &
SERVER_PID=$!

# wait for health
for i in $(seq 1 120); do
  if curl -sf http://localhost:$PORT/health > /dev/null 2>&1; then
    echo "server healthy after ${i}0s"
    break
  fi
  sleep 10
done
curl -sf http://localhost:$PORT/health > /dev/null || { echo "SERVER FAILED"; exit 1; }

SAMPLING=()
if [ -n "$TEMP" ]; then
  SAMPLING=(--sampling.temperature "$TEMP")
fi

OUT=$AGENTPTB_WORKSPACE/evals/$TAG
mkdir -p "$OUT"

echo "=== terminal-bench-2: $N_TB2 tasks ==="
uv run eval terminal-bench-2-v1 -m "$CKPT" \
  --client.base-url "http://127.0.0.1:$PORT/v1" \
  --env.agent.harness.id pi --env.agent.runtime.type broker \
  --env.agent.runtime.ready_timeout_seconds 1800 \
  "${SAMPLING[@]}" \
  -n "$N_TB2" -r 1 -s -c 24 --no-push --rich false \
  -o "$OUT/tb2" 2>&1 | tail -15

echo "=== swe-bench-verified: $N_SWE tasks ==="
uv run eval swebench-verified-v1 -m "$CKPT" \
  --client.base-url "http://127.0.0.1:$PORT/v1" \
  --env.agent.harness.id pi --env.agent.runtime.type broker \
  --env.agent.runtime.ready_timeout_seconds 1800 \
  "${SAMPLING[@]}" \
  -n "$N_SWE" -r 1 -s -c 64 --no-push --rich false \
  -o "$OUT/swe" 2>&1 | tail -15

# shutdown server, free GPUs
kill $SERVER_PID 2>/dev/null || true
sleep 5
pkill -9 -f "PRIME-RL" 2>/dev/null || true

echo "=== scores ==="
uv run python $AGENTPTB_WORKSPACE/scripts/score.py "$OUT/tb2/traces.jsonl" "$OUT/swe/traces.jsonl"
echo "=== DONE $TAG ==="