#!/usr/bin/env python3
"""bench-probe.py — RUNS ON DANKSERVER. The campaign's decode probe: one
/completion request, a prompt assembled from calib.jsonl to ~TARGET tokens,
temperature 0, max_tokens 512. Reports the engine's own timings: decode rate,
prefill rate, draft accept. Run 3x; the median is the arm's number."""
import json, sys, urllib.request

TARGET = 43600  # tokens
BYTES_PER_TOK = 3.8
URL = "http://localhost:7770/v1/completions"

def build_prompt():
    want = int(TARGET * BYTES_PER_TOK)
    parts = []
    n = 0
    with open("/home/dank/corpora/out/calib.jsonl") as f:
        for line in f:
            d = json.loads(line)
            p = d.get("prompt") or d.get("messages", [{}])[0].get("content", "")
            parts.append(p)
            n += len(p.encode())
            if n >= want:
                break
    return "".join(parts)

def probe(salt):
    prompt = f"[probe {salt}]\n" + build_prompt()
    body = json.dumps({
        "model": "frank-archon",
        "prompt": prompt, "temperature": 0.0, "max_tokens": 512,
    }).encode()
    req = urllib.request.Request(URL, data=body,
        headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=1800) as r:
        d = json.load(r)
    t = d.get("timings", {})
    return {
        "prompt_n": t.get("prompt_n"),
        "prefill_tps": round(t.get("prompt_per_second", 0), 1),
        "pred_n": t.get("predicted_n"),
        "decode_tps": round(t.get("predicted_per_second", 0), 2),
        "draft_n": t.get("draft_n"),
        "accept": (round(t["draft_n_accepted"] / t["draft_n"], 4)
                   if t.get("draft_n") else None),
    }

if __name__ == "__main__":
    print(json.dumps(probe(sys.argv[1] if len(sys.argv) > 1 else "a")))