#!/usr/bin/env python3
"""Cache-busted bs=1 decode bench against a vLLM OpenAI seat + MTP acceptance."""
import json, sys, time, urllib.request, statistics, random, argparse

TOPICS = ["a coral reef ecosystem", "the Roman aqueduct system", "how lithium-ion cells degrade",
          "the history of the printing press", "how radar altimeters work", "glacier mass balance",
          "the design of the Saturn V F-1 engine", "how sourdough fermentation works"]

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", default="http://10.250.50.54:8015")
    ap.add_argument("--model", default="qwen3.8-27b-uncensored")
    ap.add_argument("--tag", default="run")
    ap.add_argument("--max-tokens", type=int, default=400)
    ap.add_argument("--out", default=None)
    ap.add_argument("--quiet", action="store_true")
    a = ap.parse_args()

    def post(p):
        req = urllib.request.Request(a.base + "/v1/chat/completions",
                                     data=json.dumps(p).encode(),
                                     headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=900) as r:
            return json.load(r)

    def spec():
        try:
            with urllib.request.urlopen(a.base + "/metrics", timeout=30) as r:
                t = r.read().decode()
        except Exception:
            return {}
        o = {}
        for k in ("vllm:spec_decode_num_draft_tokens_total",
                  "vllm:spec_decode_num_accepted_tokens_total"):
            s = 0.0
            for ln in t.splitlines():
                if ln.startswith(k) and not ln.startswith("#"):
                    try:
                        s += float(ln.rsplit(" ", 1)[1])
                    except Exception:
                        pass
            o[k] = s
        return o

    random.seed(1234)
    # warmup
    post({"model": a.model, "messages": [{"role": "user", "content": "Say hello."}],
          "max_tokens": 16, "temperature": 0})

    rates, accs = [], []
    for i, t in enumerate(TOPICS):
        nonce = random.randint(10**9, 10**10)
        prompt = f"[session {nonce}] Write a detailed technical explanation of {t}. Be thorough and specific."
        b = spec(); t0 = time.perf_counter()
        r = post({"model": a.model, "messages": [{"role": "user", "content": prompt}],
                  "max_tokens": a.max_tokens, "temperature": 0, "stream": False})
        w = time.perf_counter() - t0; af = spec()
        ct = r["usage"]["completion_tokens"]
        dd = af.get("vllm:spec_decode_num_draft_tokens_total", 0) - b.get("vllm:spec_decode_num_draft_tokens_total", 0)
        da = af.get("vllm:spec_decode_num_accepted_tokens_total", 0) - b.get("vllm:spec_decode_num_accepted_tokens_total", 0)
        ar = da / dd if dd else 0.0
        rates.append(ct / w); accs.append(ar)
        if not a.quiet:
            print(f"  {i+1}. gen={ct} {w:.2f}s -> {ct/w:6.2f} tok/s  MTP={ar*100:5.1f}%", flush=True)

    res = {"tag": a.tag, "model": a.model,
           "tok_s_median": statistics.median(rates),
           "tok_s_mean": statistics.mean(rates),
           "tok_s_min": min(rates), "tok_s_max": max(rates),
           "mtp_accept_median": statistics.median(accs),
           "rates": rates, "accs": accs}
    print(f"\n[{a.tag}] median {res['tok_s_median']:.2f} tok/s  mean {res['tok_s_mean']:.2f}  "
          f"(min {res['tok_s_min']:.2f}/max {res['tok_s_max']:.2f})  MTP {res['mtp_accept_median']*100:.1f}%")
    if a.out:
        json.dump(res, open(a.out, "w"), indent=2)
    return 0

if __name__ == "__main__":
    sys.exit(main())