#!/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 NORMAL and continues: RADIANCE_CALIB_TOKENS
  is a pooled floor and this campaign's corpus is 10.03M tokens against a 12M
  floor, so round 1 seals nothing by construction. A tap that LOSES bytes fails
  the run, and a run that never sees the engine's completion line exits non-zero.
  The first version failed on a flat round, which aborted exactly one round too
  early (2026-10-29).
- before any prompt, the tap must carry the rig's own marker
  (`--tap-marker`, written by run-qwen38-bf16.sh calib) naming THIS
  container, and that container must be running. Measured 2026-10-29:
  with neither check, four rounds of 10M tokens went into restored
  production after the rig's boot had died, and recorded 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 round_verdict(before, after, done):
    """What one round's tap bytes mean.

    'done'   the engine logged that it sealed the tap — the run is over
    'shrink' the tap LOST bytes: reset, pruned or replaced under us
    'grew'   sealed bytes this round
    'flat'   sealed nothing — which is legitimate: RADIANCE_CALIB_TOKENS is a
             pooled floor, and a round that does not reach it has nothing to
             write yet (the campaign's corpus is 10.03M tokens against a 12M
             floor, so round 1 is always flat). Failing on it aborted the run
             exactly one round too early.
    """
    if done:
        return "done"
    if after < before:
        return "shrink"
    if after == before:
        return "flat"
    return "grew"


def require_rig(a):
    """Refuse unless the rig that opened this tap is up and serving right now.

    Three facts, in order: the runner wrote its marker into the tap dir, the
    marker names the container this pusher targets, and that container is
    running. A marker without a live container is a boot that died; a live
    container without a marker is some other engine on the port — on 2026-10-29
    it was restored production, which took four rounds of 10M tokens before the
    old driver admitted it had recorded nothing.
    """
    marker = os.path.join(a.tap_dir, a.tap_marker)
    if not os.path.isfile(marker):
        qc.refuse(f"no rig marker at {marker}: this tap was not opened by "
                  "run-qwen38-bf16.sh calib, so a push here would feed whatever "
                  "answers the port and record nothing")
    try:
        lines = open(marker, encoding="utf-8", errors="replace").read().splitlines()
    except OSError as exc:
        qc.refuse(f"cannot read the rig marker {marker}: {exc}")
    meta = dict(line.split("=", 1) for line in lines if "=" in line)
    who = meta.get("container", "")
    if who != a.container:
        qc.refuse(f"the rig marker names container {who!r} and this pusher targets "
                  f"{a.container!r} — wrong tap, or a stale marker")
    try:
        p = subprocess.run(["podman", "inspect", "-f", "{{.State.Running}}", a.container],
                           capture_output=True, text=True, timeout=30)
    except Exception as exc:
        qc.refuse(f"cannot inspect the rig container {a.container}: {exc}")
    if p.returncode != 0 or p.stdout.strip() != "true":
        qc.refuse(f"the rig container {a.container} is not running — whatever is "
                  "answering the port, it is not the calibration rig")


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("--tap-marker", default=".rig-boot",
                    help="the file run-qwen38-bf16.sh calib writes into the tap dir; without it the pusher refuses, because an unmarked tap is not a rig")
    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)
    require_rig(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)
        verdict = round_verdict(before, after, False)
        if verdict == "shrink":
            print(f"FAIL: the tap at {a.tap_dir} LOST bytes this round "
                  f"({before} B before, {after} B after) — it was reset or replaced "
                  "under us", flush=True)
            sys.exit(1)
        if verdict == "flat":
            print(f"round {rnd}: the tap sealed nothing yet (byte floor not reached; "
                  f"{before} B) — continuing", flush=True)
    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()