#!/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).
- a round that seals no tap bytes is a FAILURE, not a retry: the engine
  answering 200 while the tap directory is wrong, full or read-only is
  how a corpus gets pushed four times for nothing.

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 os
import subprocess
import sys
import time
import urllib.request

_HERE = os.path.dirname(os.path.abspath(__file__)) if "__file__" in globals() else os.getcwd()
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)

import qwen38_common as qc  # noqa: E402 - sibling module, same directory

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, model):
    body = json.dumps({
        "model": model,
        "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 push_round(url, prompts, model, parallel):
    """Push the whole corpus once; returns how many requests were accepted."""
    ok = 0
    with concurrent.futures.ThreadPoolExecutor(parallel) as ex:
        for good in ex.map(lambda p: push_one(url, p, model), prompts):
            ok += 1 if good else 0
    return ok


def tap_bytes(path):
    """Bytes the tap has sealed so far; 0 when the directory is absent.

    A tap whose directory is wrong, full or read-only stops writing while the
    engine keeps answering 200 — the driver then re-pushes the whole corpus for
    every round it is allowed and the loss is invisible until the hours are gone.
    """
    total = 0
    try:
        for root, _, files in os.walk(path):
            for f in files:
                try:
                    total += os.path.getsize(os.path.join(root, f))
                except OSError:
                    pass
    except OSError:
        return 0
    return total


def parse_args():
    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("--model", default="frank-archon")
    ap.add_argument("--tap-dir", default="/home/dank/campaign/qwen38-w5nl64/calib")
    ap.add_argument("--corpus", default="/home/dank/corpora/out/calib.jsonl")
    ap.add_argument("--max-rounds", type=int, default=4)
    return ap.parse_args()


def load_prompts(a) -> list[str]:
    if not os.path.isfile(a.corpus):
        qc.refuse(f"no corpus at {a.corpus}")
    prompts = []
    for lineno, rec in qc.iter_jsonl(a.corpus):
        if not isinstance(rec.get("prompt"), str):
            qc.refuse(f"{a.corpus}:{lineno} has no string 'prompt' field")
        prompts.append(rec["prompt"])
    if not prompts:
        qc.refuse(f"the corpus at {a.corpus} holds no prompts")
    return prompts


def main():
    a = parse_args()
    url = f"http://127.0.0.1:{a.port}/v1/completions"
    prompts = load_prompts(a)
    chars = sum(len(p) for p in prompts)
    print(f"corpus: {len(prompts)} prompts, ~{chars // 4} tokens (chars/4)",
          flush=True)

    for rnd in range(1, a.max_rounds + 1):
        if tap_done(a.container):
            break
        before = tap_bytes(a.tap_dir)
        t0 = time.time()
        ok = push_round(url, prompts, a.model, a.parallel)
        dt = max(time.time() - t0, 1e-9)
        print(f"round {rnd}: {ok}/{len(prompts)} ok in {dt:.0f}s "
              f"({chars / 4 / dt:.0f} tokens/s est)", flush=True)
        if ok == 0:
            print("FAIL: not one prompt was accepted; fix the engine or the argv "
                  "before spending another round", flush=True)
            sys.exit(1)
        if tap_done(a.container):
            break
        after = tap_bytes(a.tap_dir)
        if after <= before:
            print(f"FAIL: the tap at {a.tap_dir} sealed nothing this round "
                  f"({before} B before, {after} B after) — the engine is serving "
                  "without recording (wrong dir, full disk, read-only mount)", flush=True)
            sys.exit(1)
    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()