#!/usr/bin/env python3
"""deep-probe.py — one chat completion at ~43.6k prompt tokens, temp 0, 512
generated, streamed; decode = completion_tokens / (total - ttft)."""
import json, time, urllib.request

TARGET_BYTES = int(43600 * 3.8)

def build():
    parts, n = [], 0
    with open("/home/dank/corpora/out/calib.jsonl") as f:
        for line in f:
            p = json.loads(line)["prompt"]
            parts.append(p)
            n += len(p.encode())
            if n >= TARGET_BYTES:
                break
    return "".join(parts)

def main(salt):
    body = json.dumps({
        "model": "frank-archon",
        "messages": [{"role": "user", "content": f"[probe {salt}] Continue and summarize this material in English, 400 words.\n\n" + build()}],
        "temperature": 0.0, "max_tokens": 512,
        "stream": True, "stream_options": {"include_usage": True},
    }).encode()
    req = urllib.request.Request("http://localhost:7770/v1/chat/completions",
        data=body, headers={"Content-Type": "application/json"})
    t0 = time.time(); ttft = None; ctok = 0; usage = {}
    with urllib.request.urlopen(req, timeout=2400) as r:
        for raw in r:
            line = raw.decode().strip()
            if not line.startswith("data:"):
                continue
            d = line[5:].strip()
            if d == "[DONE]":
                break
            j = json.loads(d)
            if j.get("usage"):
                usage = j["usage"]
            for ch in j.get("choices", []):
                delta = (ch.get("delta") or {})
                piece = delta.get("content") or delta.get("reasoning_content") or ""
                if piece and ttft is None:
                    ttft = time.time() - t0
                ctok += 1
    total = time.time() - t0
    gen = usage.get("completion_tokens") or ctok
    dec = gen / (total - (ttft or 0)) if total > (ttft or 0) else 0
    print(json.dumps({"salt": salt, "prompt_tokens": usage.get("prompt_tokens"),
                      "completion_tokens": gen, "ttft_s": round(ttft or 0, 2),
                      "total_s": round(total, 1), "decode_tps": round(dec, 2)}))

if __name__ == "__main__":
    import sys
    main(sys.argv[1] if len(sys.argv) > 1 else "x")