#!/usr/bin/env python3
"""Benchmark llama.cpp server: prefill (pp) and generation (tg) throughput."""
import json, sys, urllib.request, argparse, random, string

def post(port, path, payload, timeout=1800):
    req = urllib.request.Request(
        f"http://localhost:{port}{path}",
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.load(r)

def make_prompt(approx_tokens):
    # ~0.75 words/token for english-ish filler; use varied words to avoid trivial cache hits
    words = ["system", "kernel", "vector", "matrix", "gradient", "tensor", "buffer",
             "context", "attention", "expert", "router", "layer", "cache", "token",
             "inference", "throughput", "latency", "quantize", "offload", "memory"]
    rnd = random.Random(1234)
    n_words = int(approx_tokens * 0.75)
    return " ".join(rnd.choice(words) for _ in range(n_words))

def bench(port, prompt_tokens, n_predict, label):
    prompt = make_prompt(prompt_tokens)
    r = post(port, "/completion", {
        "prompt": prompt,
        "n_predict": n_predict,
        "temperature": 0.7,
        "cache_prompt": False,
    })
    t = r.get("timings", {})
    pp_n = t.get("prompt_n", 0); pp_ms = t.get("prompt_ms", 0)
    tg_n = t.get("predicted_n", 0); tg_ms = t.get("predicted_ms", 0)
    pp_tps = t.get("prompt_per_second") or 0
    tg_tps = t.get("predicted_per_second") or 0
    print(f"{label:<28} pp: {pp_n:>7} tok @ {pp_tps:>8.1f} t/s | tg: {tg_n:>4} tok @ {tg_tps:>6.2f} t/s",
          flush=True)
    return {"label": label, "pp_n": pp_n, "pp_tps": pp_tps, "tg_n": tg_n, "tg_tps": tg_tps,
            "pp_ms": pp_ms, "tg_ms": tg_ms}

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=18010)
    ap.add_argument("--label", default="run")
    ap.add_argument("--depths", default="512,8192,32768,131072")
    ap.add_argument("--n-predict", type=int, default=128)
    ap.add_argument("--json-out", default=None)
    a = ap.parse_args()

    # Warm up: the first request after load pays one-off CUDA graph / kernel
    # setup that otherwise shows up as a bogus ~30% slower first data point.
    try:
        post(a.port, "/completion", {"prompt": make_prompt(64), "n_predict": 16,
                                     "temperature": 0, "cache_prompt": False})
    except Exception as e:
        print(f"warmup failed: {e}")

    results = []
    for d in [int(x) for x in a.depths.split(",")]:
        try:
            results.append(bench(a.port, d, a.n_predict, f"{a.label} @{d//1024}k" if d >= 1024 else f"{a.label} @{d}"))
        except Exception as e:
            print(f"{a.label} @{d}: FAILED {type(e).__name__}: {e}")
    if a.json_out:
        with open(a.json_out, "w") as f:
            json.dump(results, f, indent=2)