#!/usr/bin/env python3
"""exp15/Q4a — paired 2x2 generator (B-hybrid, orchestrator-ratified 2026-07-18).

Generates the full square {flash, pro} x {raw draft, +glm-editor}, x2 draft seeds, on the 7 exp14
meaning-trap chunks. PAIR INTEGRITY IS SUPREME (orchestrator pin): both flash and pro drafts get
BYTE-IDENTICAL prompts + injection — the exp14 injected_ids reconstruction (bilingual_glossary for
the draft, editor_constraint for the editor), imperfect (no sticky) but SHARED by both sides; we do
NOT "improve" it with the S2' mem_select port. The ONLY variable in the draft row is the draft model.

Assembly (exp14 harness, matches the saved flash-draft provenance):
  draft  : messages_with_injection(translator_system(), bilingual_glossary, translator_user(src))
           deepseek-v4-{flash|pro}, temp 0.3, thinking-ON (never disabled — echo-mine mandate)
  editor : messages_with_injection(editor_system("discourse"), editor_constraint, editor_user(src,draft))
           glm-5, temp 0.4, thinking-OFF  (same editor for both sides -> editor-delta is model-clean)

Why regenerate the flash side too (vs reuse on-disk r["draft"]/arm C): the on-disk flash came from
the Go backend (real sticky injection); regenerating flash with the SAME reconstruction assembly as
pro GUARANTEES the flash<->pro pairing the orchestrator ranks above reuse. On-disk backend flash is
retained (records.json / exp14b arm C) as a provenance cross-check, DET-scored in q4a_score.py.

Valley guard: deepseek-v4-pro calls are BLOCKED in DeepSeek peak windows (UTC 01-04 & 06-10, the
x2 surge) — abort with a defer message; flash allowed anytime. Per-call predicted-cost gate + own
ledger (every attempt incl. failures persisted, D30.10). Resumable (skip non-empty existing files).
"""
from __future__ import annotations
import argparse
import json
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "exp14"))
import exp14_common as C
import exp14_prompts as P

SD = Path("/home/ubuntu/books/gu-zhenren/exp15")
OUT = SD / "q4a"
LEDGER = SD / "q4a_gen_costs.jsonl"
RECS = {(r["chapter"], r["chunk_idx"]): r
        for r in json.load(open("/home/ubuntu/books/gu-zhenren/rerun/records.json"))}
INJ = json.load(open("/home/ubuntu/books/gu-zhenren/exp14/injection_blocks.json"))

TRAP_CHUNKS = ["6.0", "7.0", "9.1", "10.1", "16.0", "17.0", "19.0"]   # 9 meaning-traps live here
SIDES = {"flash": "deepseek-v4-flash", "pro": "deepseek-v4-pro"}
SEEDS = [0, 1, 2]
DRAFT_TEMP, EDIT_TEMP = 0.3, 0.4
# DRAFT_MAX raised 8000->16000: deepseek-v4-pro reasoning can consume the whole 8000 cap and TRUNCATE the
# translation (seen: pro_raw/7.0.s0 finish=length, reasoning starved content -> a false "omission" of the
# lost tail; adversarial verification 2026-07-18). finish=length is now REJECTED in _gen (never accepted).
DRAFT_MAX, EDIT_MAX = 16000, 12000
PEAK_UTC_HOURS = {1, 2, 3, 6, 7, 8, 9}     # DeepSeek surge windows (01-04 & 06-10 UTC) — pro forbidden
GEN_HARD_STOP = 1.10                         # cumulative gen ceiling ($; leaves room for judging under $2)

CAPS = {"deepseek-v4-flash": 0.03, "deepseek-v4-pro": 0.06, "glm-5": 0.08}


def _now_utc_hour():
    return datetime.now(timezone.utc).hour


def _src_inj(cid):
    ch, ck = cid.split("."); k = (int(ch), int(ck))
    b = INJ.get(f"{ch}/{ck}", {})
    return RECS[k]["source"], b.get("bilingual_glossary", ""), b.get("editor_constraint", "")


def _ledger_total():
    if not LEDGER.exists():
        return 0.0
    t = 0.0
    for line in LEDGER.read_text(encoding="utf-8").splitlines():
        try:
            t += json.loads(line).get("cost", 0.0)
        except Exception:
            pass
    return t


def _gen(sp, model, messages, temp, max_out, meta):
    """Gate -> call -> persist EVERY attempt (incl. gate-block/error). Returns text or None."""
    in_est = sum(len(m["content"]) for m in messages) // 2
    ok, pred = sp.guard(model, in_est, max_out)
    if not ok:
        sp.record({**meta, "model": model, "keyenv": model, "usage": {},
                   "err": f"per-call gate pred ${pred:.4f} > cap ${CAPS.get(model)}", "gate_block": True})
        print(f"    [GATE BLOCK] {meta} pred ${pred:.4f}")
        return None
    total = _ledger_total()
    if total + pred > GEN_HARD_STOP:
        print(f"    [GEN HARD-STOP] {total:.3f}+{pred:.4f} > {GEN_HARD_STOP} — stopping generation")
        return "__STOP__"
    s = C.spec(model, temp=temp, max_tokens=max_out)
    text, err, usage = C.call(s, messages)
    fin = usage.get("finish_reason")
    if not err and fin == "length":     # truncated output (reasoning ate the cap) — NOT a complete draft
        err = f"truncated|finish=length|completion={usage.get('completion_tokens')}"
    sp.record({**meta, "model": model, "keyenv": model, "usage": usage, "err": err})
    if err or not text:
        print(f"    [BAD] {meta} err={err} finish={fin}")
        return None
    return text


def run(limit_chunks=0, sanity=False):
    OUT.mkdir(parents=True, exist_ok=True)
    for cell in ("flash_raw", "pro_raw", "flash_glm", "pro_glm"):
        (OUT / cell).mkdir(exist_ok=True)
    sp = C.Spender(LEDGER, CAPS)
    # sanity gate: FIRST chunk, seed 0 only, but BOTH sides -> flash+pro draft AND editor (4 cells), so
    # the pro/deepseek path (content non-empty, thinking-ON) is exercised before the full valley run.
    chunks = TRAP_CHUNKS[:1] if sanity else (TRAP_CHUNKS[:limit_chunks] if limit_chunks else TRAP_CHUNKS)
    seeds = [0] if sanity else SEEDS
    ed_sys = P.editor_system("discourse")
    tr_sys = P.translator_system()
    made, skipped = 0, 0
    for cid in chunks:
        src, gloss, edconstr = _src_inj(cid)
        tr_user = P.translator_user(src)
        for side, dmodel in SIDES.items():
            for seed in seeds:
                # ---- draft ----
                dpath = OUT / f"{side}_raw" / f"{cid}.s{seed}.txt"
                if dpath.exists() and dpath.read_text(encoding="utf-8").strip():
                    draft = dpath.read_text(encoding="utf-8"); skipped += 1
                else:
                    if dmodel == "deepseek-v4-pro" and _now_utc_hour() in PEAK_UTC_HOURS:
                        print(f"  [VALLEY DEFER] pro draft {cid} s{seed}: UTC hour {_now_utc_hour()} in peak "
                              f"{sorted(PEAK_UTC_HOURS)} — rerun in valley."); return
                    msgs = C.messages_with_injection(tr_sys, gloss, tr_user)
                    draft = _gen(sp, dmodel, msgs, DRAFT_TEMP, DRAFT_MAX,
                                 {"cell": f"{side}_raw", "chunk": cid, "seed": seed, "stage": "draft"})
                    if draft == "__STOP__":
                        return
                    if not draft:
                        continue
                    dpath.write_text(draft, encoding="utf-8"); made += 1
                    print(f"  {side}_raw {cid} s{seed}: {len(draft)}c  gen-cum ${_ledger_total():.4f}")
                # ---- editor over that draft ----
                epath = OUT / f"{side}_glm" / f"{cid}.s{seed}.txt"
                if epath.exists() and epath.read_text(encoding="utf-8").strip():
                    skipped += 1
                else:
                    emsgs = C.messages_with_injection(ed_sys, edconstr, P.editor_user(src, draft))
                    final = _gen(sp, "glm-5", emsgs, EDIT_TEMP, EDIT_MAX,
                                 {"cell": f"{side}_glm", "chunk": cid, "seed": seed, "stage": "edit"})
                    if final == "__STOP__":
                        return
                    if not final:
                        continue
                    epath.write_text(final, encoding="utf-8"); made += 1
                    print(f"  {side}_glm {cid} s{seed}: {len(final)}c  gen-cum ${_ledger_total():.4f}")
    print(f"\ngen done: {made} made, {skipped} reused. gen total ${_ledger_total():.4f}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--limit-chunks", type=int, default=0)
    ap.add_argument("--sanity", action="store_true",
                    help="gate: 1 chunk, seed 0, BOTH sides (flash+pro draft+editor = 4 cells)")
    a = ap.parse_args()
    print(f"[q4a-gen] UTC hour {_now_utc_hour()} (peaks {sorted(PEAK_UTC_HOURS)}); "
          f"gen ledger total ${_ledger_total():.4f}; hard-stop ${GEN_HARD_STOP}")
    run(a.limit_chunks, a.sanity)


if __name__ == "__main__":
    main()