#!/bin/bash
# Run on the LOCAL client. Opens the tunnel and runs term-llm's built-in
# benchmark against the 'colabvllm' provider (type: vllm) in the same shell.
#
#   ./run_tl_bench.sh quick
#   ./run_tl_bench.sh decode
#   ./run_tl_bench.sh prefill
#   ./run_tl_bench.sh balanced --input-tokens 128,512,2048,8192,12000 --output-tokens 256
#
# Same layout and environment variables as tunnel_and_termllm.sh.
set -u
MODE="${1:-quick}"
shift || true

SESSION="${COLAB_SESSION:-flashnext}"
LOCAL_PORT="${LOCAL_PORT:-18000}"
HERE="$(cd "$(dirname "$0")" && pwd)"
KEY="$HERE/id_colab"
COLAB_BIN="$(command -v colab)"

export XDG_CONFIG_HOME="$HERE/term-llm-config"
export XDG_DATA_HOME="$HERE/term-llm-data"
mkdir -p "$XDG_DATA_HOME"
cd "$HERE"

ssh -i "$KEY" -o IdentitiesOnly=yes \
  -o ProxyCommand="$COLAB_BIN ssh --proxy-mode -s $SESSION --identity $KEY" \
  -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR \
  -o ServerAliveInterval=15 -o ExitOnForwardFailure=yes \
  -N -L "$LOCAL_PORT:127.0.0.1:8000" root@colab-runtime > tunnel_bench.log 2>&1 &
TUNNEL_PID=$!
trap 'kill $TUNNEL_PID 2>/dev/null; wait $TUNNEL_PID 2>/dev/null' EXIT

for i in $(seq 1 30); do
  if curl -s -m 5 "http://127.0.0.1:$LOCAL_PORT/v1/models" >/dev/null 2>&1; then echo "tunnel ready after ${i}s"; break; fi
  sleep 1
done

START=$(date +%s.%N)
term-llm --no-session benchmark -p colabvllm --mode "$MODE" --runs 3 --warmup 1 \
  --json --jsonl "bench_tl_${MODE}.jsonl" "$@" > "bench_tl_${MODE}.json" 2> "bench_tl_${MODE}.log"
RC=$?
END=$(date +%s.%N)
echo "mode=$MODE rc=$RC ELAPSED=$(python3 -c "print(round($END-$START,2))")s"
tail -20 "bench_tl_${MODE}.log"