#!/usr/bin/env python3
"""pp-bench.py — drive the PP head API and report server-side timings.

Runs on the HEAD machine (Dankclaw). Targets the running PP head
(serve-pp-head.sh, port from argument). Builds a ~Nk-token prose prompt,
sends one warmup completion, then REPS measured completions, and reads the
authoritative `slot print_timing` lines from the head log.

Usage:
    python3 pp-bench.py <prompt_k> <out_tokens> <reps> <port> <head_log>
"""
import json
import re
import sys
import time
import urllib.request

prompt_k, out_tokens, reps, port, head_log = sys.argv[1:6]
out_tokens = int(out_tokens)
reps = int(reps)
target = int(prompt_k) * 1000

# ~38 tokens per 160-char prose unit; size the filler to hit the target.
unit = ("The careful study of natural language reveals patterns of meaning "
        "that emerge from the deliberate arrangement of words and phrases. ")
approx = max(1, target // 38)
prompt = unit * approx + (
    "\n\nNow, in a detailed and flowing paragraph, "
    "continue this essay about the philosophy of "
    "language and its influence on modern computation. "
)


def req(n):
    body = json.dumps({"prompt": prompt, "max_tokens": n,
                       "temperature": 0.0, "stream": False}).encode()
    r = urllib.request.Request(
        f"http://127.0.0.1:{port}/v1/completions",
        data=body, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(r, timeout=7200) as resp:
        return json.load(resp)


def log_timings():
    try:
        lines = open(head_log, encoding="utf-8", errors="ignore").read().splitlines()
    except OSError:
        return []
    out = []
    for l in lines:
        if re.search(r"prompt eval time|eval time|draft acceptance", l):
            out.append(l.strip())
    return out


GEN_RE = re.compile(
    r"eval time =\s+([\d.]+) ms /\s+(\d+) tokens.*?([\d.]+) tokens per second")
PRE_RE = re.compile(
    r"prompt eval time =\s+([\d.]+) ms /\s+(\d+) tokens.*?([\d.]+) tokens per second")
ACC_RE = re.compile(r"draft acceptance = ([\d.]+)")

# warmup (first run pays graph-compile / JIT costs)
warm = req(16)
print(f"warmup: prompt_tokens={warm['usage']['prompt_tokens']} "
      f"completion_tokens={warm['usage']['completion_tokens']}")

gen_rates, prefill_rates, accepts = [], [], []
for i in range(reps):
    t0 = time.time()
    out = req(out_tokens)
    dt = time.time() - t0
    u = out["usage"]
    pt, ct = u["prompt_tokens"], u["completion_tokens"]
    gen = pre = acc = None
    for l in log_timings():
        if "prompt eval" not in l:
            m = GEN_RE.search(l)
            if m:
                gen = float(m.group(3))
        m = PRE_RE.search(l)
        if m:
            pre = float(m.group(3))
        m = ACC_RE.search(l)
        if m:
            acc = float(m.group(1))
    gen_rates.append(gen if gen is not None else ct / dt)
    if pre is not None:
        prefill_rates.append(pre)
    if acc is not None:
        accepts.append(acc)
    print(f"rep {i + 1}: prompt={pt} gen={ct} wall={dt:.1f}s "
          f"gen={gen_rates[-1]:.2f} t/s "
          f"prefill={f'{pre:.1f}' if pre is not None else 'n/a'} t/s "
          f"acceptance={f'{acc:.3f}' if acc is not None else 'n/a'}")

print("== RESULT pp ==")
print(f"avg gen tok/s   = {sum(gen_rates) / len(gen_rates):.2f}")
if prefill_rates:
    print(f"avg prefill t/s = {sum(prefill_rates) / len(prefill_rates):.1f}")
if accepts:
    print(f"avg acceptance  = {sum(accepts) / len(accepts):.3f}")
for l in log_timings()[-3:]:
    print("  " + l)