#!/usr/bin/env python3
"""Decode and prefill speed of any OpenAI-compatible server, timed on the client.

  speed-test.py [--port 8082] [--label NAME] [--sizes 0,7000,27000] [--runs 2] [--max-tokens 512]

For each size: a cold prompt (a nonce first, so no cache can be reused) of about
that many tokens of source code plus a summary task, thinking off, streamed.
Prefill = prompt tokens / time to first token; decode = tokens after the first /
time after the first. Size 0 is a short coding prompt. Appends JSON lines to --out.
"""
import argparse, glob, json, os, random, time, urllib.request

ap = argparse.ArgumentParser()
ap.add_argument("--port", type=int, default=8082)
ap.add_argument("--label", default="")
ap.add_argument("--sizes", default="0,7000,27000")
ap.add_argument("--runs", type=int, default=2)
ap.add_argument("--max-tokens", type=int, default=512)
ap.add_argument("--src", default="/opt/tersler/src/llama.cpp/src")
ap.add_argument("--out", default="/opt/tersler/bench/speed.jsonl")
a = ap.parse_args()
URL = f"http://127.0.0.1:{a.port}/v1/chat/completions"

text = "\n\n".join(open(f, errors="replace").read() for f in sorted(glob.glob(os.path.join(a.src, "*.cpp"))))
SHORT = "Write a Python function that merges overlapping intervals, with a docstring and three test cases."
TASK = "\n\nSummarise what the code above does, file by file, in about 300 words."


def run(size):
    nonce = f"[{random.random():.9f}]\n"
    content = nonce + (SHORT if size == 0 else text[:size * 4] + TASK)   # ~4 characters per token of code
    body = {"messages": [{"role": "user", "content": content}], "max_tokens": a.max_tokens, "stream": True,
            "temperature": 0.6, "top_p": 0.95, "stream_options": {"include_usage": True},
            "chat_template_kwargs": {"enable_thinking": False}}
    req = urllib.request.Request(URL, json.dumps(body).encode(), {"Content-Type": "application/json"})
    t0 = time.time(); t1 = None; n = 0; usage = {}
    with urllib.request.urlopen(req, timeout=3600) as r:
        for line in r:
            line = line.strip()
            if not line.startswith(b"data:") or line == b"data: [DONE]":
                continue
            ev = json.loads(line[5:])
            usage = ev.get("usage") or usage
            for ch in ev.get("choices", []):
                d = ch.get("delta", {})
                if d.get("content") or d.get("reasoning_content"):
                    n += 1
                    t1 = t1 or time.time()
    t2 = time.time()
    pt = usage.get("prompt_tokens") or len(content) // 4
    ct = usage.get("completion_tokens") or n
    return {"size": size, "prompt_tokens": pt, "completion_tokens": ct, "ttft_s": round(t1 - t0, 2),
            "pp_tok_s": round(pt / (t1 - t0), 1), "tg_tok_s": round((ct - 1) / (t2 - t1), 1) if ct > 1 else None,
            "usage_reported": bool(usage)}


for size in [int(x) for x in a.sizes.split(",")]:
    for i in range(a.runs):
        rec = run(size)
        rec.update(label=a.label, run=i, date=time.strftime("%Y-%m-%dT%H:%M:%S"))
        print(json.dumps(rec), flush=True)
        with open(a.out, "a") as f:
            f.write(json.dumps(rec) + "\n")