"""Tiny single-stream latency probe over the OpenAI HTTP API.

Usage: python3 pp_time.py BASE_URL "PROMPT"
Prints: TTFT_ms=.. TPOT_ms=.. tok_full=.. tok_s=..

TTFT ~ wall time of a max_tokens=1 request (prefill + 1 token).
TPOT ~ (t_full - t_1) / (tokens_full - tokens_1), using the server's reported
completion_tokens so it is exact even if generation stops early.
"""
import json
import sys
import time
import urllib.request

base = sys.argv[1].rstrip("/")
prompt = sys.argv[2]


def req(max_tokens):
    body = json.dumps({
        "model": "qwen3-8b",
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": 0,
        "stream": False,
    }).encode()
    r = urllib.request.Request(base + "/v1/chat/completions", body,
                               {"Content-Type": "application/json"})
    t = time.time()
    d = json.load(urllib.request.urlopen(r, timeout=600))
    dt = time.time() - t
    ct = d.get("usage", {}).get("completion_tokens")
    return dt, ct


t1, c1 = req(1)
tF, cF = req(160)
ttft = t1 * 1000.0
denom = (cF - c1) if (cF and c1 and cF > c1) else None
if denom:
    tpot = (tF - t1) / denom * 1000.0
    print(f"TTFT_ms={ttft:.1f} TPOT_ms={tpot:.2f} tok_full={cF} tok_s={1000.0/tpot:.1f}")
else:
    print(f"TTFT_ms={ttft:.1f} TPOT_ms=nan tok_full={cF} tok_s=nan")