"""Single-stream decode throughput. Usage: bench_decode.py <port> <model> [n_tokens]"""
import json
import sys
import time
import urllib.request

PORT, MODEL = int(sys.argv[1]), sys.argv[2]
N = int(sys.argv[3]) if len(sys.argv) > 3 else 512

payload = {
    "model": MODEL,
    "messages": [
        {
            "role": "user",
            "content": "Count from 1 to 300 in words, one number per line (one, two, three, ...).",
        }
    ],
    "temperature": 0,
    "max_tokens": N,
    "ignore_eos": True,
}
req = urllib.request.Request(
    f"http://127.0.0.1:{PORT}/v1/chat/completions",
    data=json.dumps(payload).encode(),
    headers={"Content-Type": "application/json"},
)
# warmup (also primes prefix cache off: different prompt)
warm = dict(payload, max_tokens=32)
urllib.request.urlopen(
    urllib.request.Request(req.full_url, data=json.dumps(warm).encode(), headers=req.headers),
    timeout=600,
).read()
t0 = time.time()
r = json.loads(urllib.request.urlopen(req, timeout=600).read())
dt = time.time() - t0
ct = r["usage"]["completion_tokens"]
print(f"{ct} completion tokens in {dt:.2f}s = {ct / dt:.1f} tok/s (single stream, MTP on)")