#!/usr/bin/env python3
"""Prefill throughput via TTFT on long prompts (prefill-dominated)."""
import argparse, json, random, statistics, time, urllib.request

def ttft(base, model, prompt):
    req = urllib.request.Request(base + "/v1/chat/completions",
        data=json.dumps({"model": model, "messages": [{"role": "user", "content": prompt}],
                         "max_tokens": 16, "temperature": 0, "stream": True,
                         "chat_template_kwargs": {"enable_thinking": False}}).encode(),
        headers={"Content-Type": "application/json"})
    t0 = time.perf_counter()
    with urllib.request.urlopen(req, timeout=900) as r:
        for line in r:
            if line.startswith(b"data: ") and b"[DONE]" not in line:
                dl = json.loads(line[6:])["choices"][0].get("delta", {})
                if dl.get("content") or dl.get("reasoning"):
                    return time.perf_counter() - t0
    return None

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", required=True)
    ap.add_argument("--model", required=True)
    ap.add_argument("--tag", default="run")
    ap.add_argument("--out", default=None)
    a = ap.parse_args()
    # NO fixed seed: a seeded nonce reproduces the previous run's prompts verbatim,
    # so --enable-prefix-caching serves them from cache and you measure cache-hit
    # TTFT, not prefill. (First attempt read 41k tok/s that way.) Lengths stay
    # identical across builds; only the nonce differs.
    rnd = random.SystemRandom()
    res = {"tag": a.tag}
    body = ""
    for reps, label in [(4000, "~6.5k"), (16000, "~26k")]:
        ts = []
        for _ in range(3):
            n = rnd.randint(10**12, 10**13)
            body = f"[doc {n}] The archived engineering report details subsystem telemetry. " * (reps // 12)
            v = ttft(a.base, a.model, body + "\n\nSummarize in one sentence.")
            if v:
                ts.append(v)
        tok = len(body) // 4
        med = statistics.median(ts)
        rate = tok / med
        res[label] = {"approx_prompt_tokens": tok, "ttft_median_s": med, "prefill_tok_s": rate}
        print(f"  {label} prompt (~{tok} tok): TTFT {med:.3f}s -> ~{rate:,.0f} tok/s prefill")
    if a.out:
        json.dump(res, open(a.out, "w"), indent=2)
    return 0

if __name__ == "__main__":
    raise SystemExit(main())