#!/usr/bin/env bash
# Benchmark the vLLM server: time-to-first-token and decode tok/s from a streaming
# /v1/chat/completions response (usage included in the final chunk).
#
# Usage: ./benchmark.sh [num_predict]
#   BASE_URL=http://localhost:18080   # direct; use :11433 to go through the proxy
set -euo pipefail

BASE_URL="${BASE_URL:-http://localhost:18080}"
MODEL="${MODEL:-foundation}"
NUM_PREDICT="${1:-300}"
PROMPT="Write a comprehensive explanation of how quantum computing works, including the principles of superposition, entanglement, and quantum gates. Provide examples of real-world applications and discuss the current state of quantum hardware."

if ! curl -sf -m 10 "$BASE_URL/v1/models" >/dev/null; then
    echo "ERROR: cannot reach $BASE_URL/v1/models — is the stack up? (docker compose up -d)" >&2
    exit 1
fi

python3 - "$BASE_URL" "$MODEL" "$PROMPT" "$NUM_PREDICT" <<'PY'
import json, sys, time, urllib.request

base, model, prompt, n = sys.argv[1], sys.argv[2], sys.argv[3], int(sys.argv[4])

def chat(stream, max_tokens):
    body = json.dumps({
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "stream": stream,
        "stream_options": {"include_usage": True} if stream else None,
    }).encode()
    req = urllib.request.Request(base + "/v1/chat/completions", data=body,
                                 headers={"Content-Type": "application/json",
                                          "Authorization": "Bearer x"})
    return urllib.request.urlopen(req, timeout=1800)

# warmup (also loads/compiles if the server is fresh)
t0 = time.monotonic()
with chat(False, 16) as r:
    json.load(r)
print(f"warmup done in {time.monotonic() - t0:.1f}s")

t0 = time.monotonic()
t_first = None
usage = {}
answer = ""
with chat(True, n) as r:
    for raw in r:
        line = raw.decode().strip()
        if not line.startswith("data: "):
            continue
        payload = line[6:]
        if payload == "[DONE]":
            break
        chunk = json.loads(payload)
        if chunk.get("usage"):
            usage = chunk["usage"]
        for ch in chunk.get("choices", []):
            delta = ch.get("delta") or {}
            # vLLM 0.27 puts the reasoning trace in `reasoning` (older builds:
            # `reasoning_content`); either is the first generated token.
            piece = (delta.get("content") or delta.get("reasoning")
                     or delta.get("reasoning_content") or "")
            if piece:
                if t_first is None:
                    t_first = time.monotonic()
                answer += piece
t_end = time.monotonic()

pt = usage.get("prompt_tokens", 0)
ct = usage.get("completion_tokens", 0)
rt = (usage.get("completion_tokens_details") or {}).get("reasoning_tokens", 0) or 0
ttft = (t_first - t0) if t_first else float("nan")
decode_window = (t_end - t_first) if t_first else 0.0

print(f"prompt tokens     : {pt}")
print(f"completion tokens : {ct} (reasoning: {rt})")
print(f"TTFT              : {ttft:.2f} s")
print(f"wall              : {t_end - t0:.2f} s")
if decode_window > 0 and ct > 1:
    print(f"decode            : {(ct - 1) / decode_window:.2f} tok/s")
print(f"sample            : {answer[:140].replace(chr(10), ' ')!r}")
PY