#!/usr/bin/env python3
"""autoresearch-bench.py — deterministic greedy battery for frank-dragoon.

Runs ON Dankclaw (no LAN jitter) against the candidate container on :7770.
Fixed embedded corpus, temperature 0, streaming so decode speed is measured
between real token arrivals (not request totals). Median of REPS samples per
metric is the reported value; medians are far more stable than means under
GPU thermal drift.

Metrics emitted (stdout, one per line):
  METRIC decode_tokens_per_sec=<float>   primary, higher better
  METRIC prefill_tokens_per_sec=<float>  secondary
  METRIC ttft_ms=<float>                 secondary, lower better
Exit non-zero if the endpoint is unhealthy or the served alias is wrong.
"""
import argparse
import json
import statistics
import sys
import time
import urllib.request

# Fixed corpus: deterministic, ~28k tokens after tokenization (same scale as
# the 2026-10-01 bench window). Built by repeating a fixed paragraph block a
# fixed number of times - NO random, NO clock, NO network.
_BLOCK = (
    "The rendering pipeline walks the scene graph once per frame, resolving "
    "material bindings into descriptor sets before any draw call is recorded. "
    "Because the allocator recycles command buffers only after their fence "
    "signals, a submission that outlives its frame must pin the pool entry; "
    "otherwise the next frame overwrites recorded commands in flight. The "
    "validation layer flags exactly this class of lifetime bug, but only for "
    "buffers it can see - GPU-side timestamps and queue-family ownership "
    "transfers hide in the driver's own bookkeeping, which is why release "
    "builds must lean on explicit barriers and why a frame-time regression "
    "often shows up as corruption one frame later than its cause.\n"
)
LONG_CTX = _BLOCK * 220  # ~28k-token prefill, matching the deployed workload
_LONG = LONG_CTX + "\nNow, in one focused paragraph (80-150 words), explain " \
    "which lifetime bugs this pipeline can still suffer and what explicit " \
    "mitigations the text implies."
_SHORT = "Design a rate limiter for a homelab reverse proxy that must not " \
    "lock out the LAN. Think it through, then give the algorithm."

REPS = 3
MAX_TOKENS = 256
REQ_TIMEOUT = 900


def stream_completion(url: str, model: str, prompt: str):
    """Greedy streaming completion.

    Returns (ttft_s, decode_tps, prefill_tps, prompt_tokens, n_tokens).
    ttft is RELATIVE to request send: time.perf_counter()'s epoch is host
    boot, and returning the raw stamp read as an 11-day first token.
    """
    body = json.dumps({
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": MAX_TOKENS,
        "temperature": 0,
        "stream": True,
        "stream_options": {"include_usage": True},
    }).encode()
    req = urllib.request.Request(
        url.rstrip("/") + "/v1/chat/completions", data=body,
        headers={"Content-Type": "application/json"})
    t0 = time.perf_counter()
    first = None
    stamps = []
    usage = {}
    with urllib.request.urlopen(req, timeout=REQ_TIMEOUT) as resp:
        for raw in resp:
            line = raw.decode(errors="replace").strip()
            if not line.startswith("data:"):
                continue
            payload = line[5:].strip()
            if payload == "[DONE]":
                break
            chunk = json.loads(payload)
            if chunk.get("usage"):
                usage = chunk["usage"]
            for ch in chunk.get("choices", []):
                d = ch.get("delta") or {}
                # reasoning AND content deltas both arrive token by token;
                # count either so greedy text length drives the rate.
                # (llama.cpp names the thinking field reasoning_content.)
                if d.get("content") or d.get("reasoning") or d.get("reasoning_content"):
                    now = time.perf_counter()
                    if first is None:
                        first = now
                    stamps.append(now)
    if first is None or len(stamps) < 2:
        # A zero-argmax degenerate answer would poison the median; raise.
        raise RuntimeError(f"stream produced {len(stamps)} tokens")
    ttft_s = first - t0
    decode_s = stamps[-1] - stamps[0]
    decode_tps = (len(stamps) - 1) / decode_s if decode_s > 0 else 0.0
    prompt_tokens = usage.get("prompt_tokens") or 0
    prefill_tps = prompt_tokens / ttft_s if prompt_tokens else 0.0
    return ttft_s, decode_tps, prefill_tps, prompt_tokens, len(stamps)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", default="http://127.0.0.1:7770")
    ap.add_argument("--model", default="frank-dragoon")
    args = ap.parse_args()

    with urllib.request.urlopen(args.url.rstrip("/") + "/v1/models",
                                timeout=10) as r:
        ids = [m["id"] for m in json.load(r)["data"]]
    if args.model not in ids:
        print(f"FATAL: served ids {ids}, expected {args.model}",
              file=sys.stderr)
        sys.exit(2)

    # Warmup (compile/graph warm + allocator settle), NOT measured.
    stream_completion(args.url, args.model, _SHORT)

    decodes, prefills, ttfts = [], [], []
    for rep in range(REPS):
        for name, prompt in (("long", _LONG), ("short", _SHORT)):
            ttft, dt, pt, ptk, ntok = stream_completion(
                args.url, args.model, prompt)
            decodes.append(dt)
            ttfts.append(ttft * 1000.0)
            if name == "long":
                prefills.append(pt)
            print(f"  rep{rep} {name:5s} ttft={ttft*1000:8.1f}ms "
                  f"prefill={pt:8.1f} t/s decode={dt:6.2f} t/s "
                  f"({ntok} tok, prompt {ptk})", file=sys.stderr)

    med = lambda xs: statistics.median(xs)
    print(f"METRIC decode_tokens_per_sec={med(decodes):.2f}")
    print(f"METRIC prefill_tokens_per_sec={med(prefills):.1f}")
    print(f"METRIC ttft_ms={med(ttfts):.1f}")


if __name__ == "__main__":
    main()