#!/usr/bin/env python3
"""calib-push.py — drive a calibration corpus through a radiance scratch
server's prefill so its tap writes layer Hessians (docs/MOE-W4.md).

Runs ON DANKSERVER next to the scratch server (the one autoresearch- /
serve-style scratch container booted with RADIANCE_CALIB_DIR set).
Reads the corpus JSONL produced by make-calib-corpus.py
({"prompt": "..."} a line), POSTs each prompt to /v1/completions with
max_tokens=1 (prefill-only cost, one token out), at N-way concurrency
matching the scratch server's --max-num-seqs budget.

Corpus rounds repeat until the engine prints the tap's completion line
("calibration wrote layer Hessians") in `podman logs <container>` — that
is the ground truth that RADIANCE_CALIB_TOKENS was reached. After that
line, further prompts are pure waste, so the driver exits.

Mechanics notes:
- requests carry no sampling params that matter (max_tokens=1, greedy).
- a per-request timeout of 30 min: at ~1k tok/s prefill a 12k-token
  prompt is a few seconds; the timeout only catches a wedged engine.
- failures print and continue: the tap pools over tokens, a few lost
  prompts do not bias the Hessians (the corpus is drawn round-robin over
  source classes, so loss is class-neutral as long as failures are rare;
  a systematic 4xx is loud in the log and must be fixed, not ignored).

Usage: calib-push.py [--port 7770] [--parallel 64] [--container arbench]
       [--corpus /home/dank/corpora/out/calib.jsonl] [--max-rounds 4]
"""
import argparse
import concurrent.futures
import json
import subprocess
import sys
import time
import urllib.request

DONE_MARK = "calibration wrote"


def tap_done(container):
    """True if the engine logged the tap's completion line."""
    try:
        out = subprocess.run(
            ["podman", "logs", "--tail", "50", container],
            capture_output=True, text=True, timeout=30,
        )
    except Exception:
        return False
    return DONE_MARK in (out.stdout + out.stderr)


def push_one(url, prompt):
    body = json.dumps({
        "model": "frank-archon",
        "prompt": prompt,
        "max_tokens": 1,
        "temperature": 0,
        "stream": False,
    }).encode()
    req = urllib.request.Request(
        url, data=body, headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=1800) as r:
            r.read()
        return True
    except Exception as e:
        print(f"FAIL {str(e)[:120]}", flush=True)
        return False


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=7770)
    ap.add_argument("--parallel", type=int, default=64)
    ap.add_argument("--container", default="arbench")
    ap.add_argument("--corpus", default="/home/dank/corpora/out/calib.jsonl")
    ap.add_argument("--max-rounds", type=int, default=4)
    a = ap.parse_args()
    url = f"http://127.0.0.1:{a.port}/v1/completions"

    prompts = []
    with open(a.corpus) as f:
        for line in f:
            line = line.strip()
            if line:
                prompts.append(json.loads(line)["prompt"])
    print(f"corpus: {len(prompts)} prompts", flush=True)

    for rnd in range(1, a.max_rounds + 1):
        if tap_done(a.container):
            break
        t0 = time.time()
        ok = 0
        with concurrent.futures.ThreadPoolExecutor(a.parallel) as ex:
            for good in ex.map(lambda p: push_one(url, p), prompts):
                ok += 1 if good else 0
        dt = time.time() - t0
        print(f"round {rnd}: {ok}/{len(prompts)} ok in {dt:.0f}s "
              f"({sum(len(p) for p in prompts)/4/dt:.0f} chars/s est)",
              flush=True)
        if tap_done(a.container):
            break
    else:
        print("max rounds exhausted, tap not done", flush=True)
        sys.exit(1)
    print("TAP_DONE: calibration wrote layer Hessians", flush=True)


if __name__ == "__main__":
    main()