"""Single-stream decode benchmark against a running vLLM server (one request at a time, wall clock per request
incl. TTFT, temperature 0, 1024 output tokens, reasoning_effort=low, 24 prompts = code/math/chat/tool x 6 from
prompts.json[40:46], 2 warm-up requests). tau (mean acceptance length incl. the bonus token) and per-position
acceptance come from the server's /metrics spec-decode counters.
usage: bench_mtp_client.py --model SERVED --out OUT.json [--port 8000]"""
import argparse
import json
import re
import time
import urllib.request

from openai import OpenAI


def metrics(port):
    txt = urllib.request.urlopen(f"http://localhost:{port}/metrics").read().decode()
    m = {"drafts": 0.0, "accepted": 0.0, "per_pos": {}}
    for line in txt.splitlines():
        if line.startswith("vllm:spec_decode_num_drafts_total"):
            m["drafts"] += float(line.split()[-1])
        elif line.startswith("vllm:spec_decode_num_accepted_tokens_total"):
            m["accepted"] += float(line.split()[-1])
        elif line.startswith("vllm:spec_decode_num_accepted_tokens_per_pos_total"):
            pos = int(re.search(r'position="(\d+)"', line).group(1))
            m["per_pos"][pos] = m["per_pos"].get(pos, 0.0) + float(line.split()[-1])
    return m


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model", required=True)
    p.add_argument("--port", default="8000")
    p.add_argument("--per-cat", type=int, default=6)
    p.add_argument("--max-tokens", type=int, default=1024)
    p.add_argument("--out", required=True)
    a = p.parse_args()
    c = OpenAI(base_url=f"http://localhost:{a.port}/v1", api_key="x", timeout=3600)
    prompts = json.load(open("/root/specexp/prompts.json"))
    by = {}
    for x in prompts:
        by.setdefault(x["cat"], []).append(x)
    sel = [x for cat in ["code", "math", "chat", "tool"] for x in by[cat][40:40 + a.per_cat]]
    kw = dict(model=a.model, max_tokens=a.max_tokens, temperature=0.0,
              extra_body={"chat_template_kwargs": {"reasoning_effort": "low"}, "ignore_eos": False})
    for x in by["chat"][:2]:
        c.chat.completions.create(messages=x["messages"], **kw)
    m0 = metrics(a.port)
    res = []
    for x in sel:
        t = time.perf_counter()
        r = c.chat.completions.create(messages=x["messages"], **kw)
        dt = time.perf_counter() - t
        msg = r.choices[0].message
        text = (getattr(msg, "reasoning_content", None) or getattr(msg, "reasoning", None) or "") + "\x00" + (msg.content or "")
        res.append({"id": x["id"], "cat": x["cat"], "n": r.usage.completion_tokens, "s": dt, "text": text})
    m1 = metrics(a.port)
    out = {"model": a.model, "per_request": res}
    for cat in ["code", "math", "chat", "tool", "all"]:
        rr = [r for r in res if cat == "all" or r["cat"] == cat]
        out[cat] = sum(r["n"] for r in rr) / sum(r["s"] for r in rr)
    d = m1["drafts"] - m0["drafts"]
    if d > 0:
        out["tau"] = 1 + (m1["accepted"] - m0["accepted"]) / d
        out["accept_per_pos"] = [(m1["per_pos"][k] - m0["per_pos"].get(k, 0.0)) / d for k in sorted(m1["per_pos"])]
    print("RESULT", json.dumps({k: (round(v, 3) if isinstance(v, float) else
                                    ([round(z, 3) for z in v] if isinstance(v, list) else v))
                                for k, v in out.items() if k != "per_request"}), flush=True)
    json.dump(out, open(a.out, "w"))


if __name__ == "__main__":
    main()