#!/usr/bin/env python3
"""exp15 — Q3 cohesion arm-runner (eval-rig; design §B3/§D2). The backend can't do carryover/lookahead
(dead knobs), so Q3 arms are eval-rigs that reuse A0's EXACT prompts + validated mem_select injection,
adding ONLY the cohesion context. To keep the contrast carryover-ONLY (no backend-vs-rig confound), the
rig generates BOTH a rig-A0 baseline (no context) and the Q3 arms; rig-A0 is validated to fall within the
backend flat-A0 <-> flat-A0' noise floor.

Modes (§B3 / §D2 Q3a/Q3a'/Q3b):
  a0        : rig baseline, no cohesion context (bridge to backend flat-A0)
  q3a-src   : + carryover = last K=3 SOURCE sentences of prev chunk (parallelism-safe)
  q3a-draft : + carryover = prev src tail ∥ prev chunk DRAFT tail (what the wave model gives)
  q3a-final : + carryover = prev src tail ∥ prev chunk FINAL tail (target-side, sequential pipeline)
  q3a-state : q3a-final + deterministic scene-state (active entities + their dst + gender from seed)
  q3b-look  : + lookahead = first M=2 SOURCE sentences of NEXT chunk (+ seam-dedup guard in scoring)

Prompt assembly mirrors render.go: system(tpl.System + "\n\n" + FewShot) [CacheBoundary] -> system(injection)
-> [system(cohesion context, read-only)] -> user(filled). Draft deepseek-v4-flash temp0.3 thinking-ON;
editor glm-5 temp0.4 thinking-OFF. Per-call gate + usage/cost ledger (D30.10); resumable (skip existing).
"""
from __future__ import annotations
import argparse
import json
import re
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import exp15_llm as L
from mem_select import (build_bank, render_glossary_block, render_editor_constraint_block, union_sticky)
from chunker import greedy_chunks, split_source_sentences

BACKEND = Path("/home/ubuntu/projects/textmachine/backend")
S2 = Path("/home/ubuntu/books/gu-zhenren/exp15/S2prime_big_flat.txt")   # 30-节 flat slice
SEED = "/home/ubuntu/books/gu-zhenren/guzhenren-seed-v2.yaml"
OUTDIR = Path("/home/ubuntu/books/gu-zhenren/exp15/arms")
# Q3 baseline + carryover source = the REAL backend flat-A0-30 (not a rig-a0). Contrast = carryover-only.
BASE_DB = "/home/ubuntu/books/gu-zhenren/exp15/anchors/flat_a0_30.db"
BUDGET_TOK = 800
K_CARRY = 3
M_LOOK = 2

BRIEF = {"book_id": "exp15", "title": "蛊真人", "source_lang": "zh", "target_lang": "ru",
         "genre": "вебновелла", "audience": "взрослые читатели вебновелл", "venuti": "0.60",
         "honorifics": "keep", "transcription": "palladius", "footnotes": "minimal"}

CTX_HEADER = ("КОНТЕКСТ ДЛЯ СВЯЗНОСТИ (предыдущий/следующий фрагмент, ТОЛЬКО ДЛЯ ПОНИМАНИЯ переходов "
              "и разрешения местоимений/времён/рода — его НЕ переводить и НЕ включать в вывод):")
STATE_HEADER = "АКТИВНЫЕ СУЩНОСТИ СЦЕНЫ (канонический перевод и род — для согласования местоимений/времён):"


def parse_template(path):
    raw = Path(path).read_text(encoding="utf-8")
    parts = raw.split("---USER---", 1)
    head, user = parts[0], parts[1].strip()
    hp = head.split("---FEWSHOT---", 1)
    system = hp[0].strip()
    fewshot = hp[1].strip() if len(hp) == 2 else ""
    return {"system": system, "fewshot": fewshot, "user": user}


def fill(tpl_text, **extra):
    vals = dict(BRIEF)
    vals.update(extra)
    out, rest = [], tpl_text
    while True:
        i = rest.find("{{")
        if i < 0:
            out.append(rest); break
        out.append(rest[:i]); rest = rest[i:]
        e = rest.find("}}")
        name = rest[2:e]
        out.append(vals[name]); rest = rest[e + 2:]
    return "".join(out)


def system_msg(tpl):
    return tpl["system"] + ("\n\n" + tpl["fewshot"] if tpl["fewshot"] else "")


def seed_entities():
    import yaml
    d = yaml.safe_load(Path(SEED).read_text(encoding="utf-8"))
    return d["terms"]


def scene_state_block(chunk_src, terms):
    """Deterministic scene-state: seed entities present in the chunk, with dst + gender (§B3 state slot)."""
    lines = []
    seen = set()
    for t in terms:
        src = t.get("src"); dst = (t.get("dst") or "").strip()
        if not src or not dst:
            continue
        variants = [src] + [a.get("alias") for a in (t.get("aliases") or []) if a.get("alias")]
        if not any(v in chunk_src for v in variants):
            continue
        if t.get("type") not in ("name", "nickname"):
            continue
        g = t.get("gender") or ""
        gtag = {"male": "муж.", "female": "жен.", "hidden": "?", "": ""}.get(g, "")
        key = dst
        if key in seen:
            continue
        seen.add(key)
        lines.append(f"- {dst}" + (f" ({gtag} род)" if gtag else ""))
    return (STATE_HEADER + "\n" + "\n".join(lines)) if lines else ""


def tail_sentences(text, k):
    s = split_source_sentences(text)
    return "".join(s[-k:]).strip() if s else ""


def head_sentences(text, m):
    s = split_source_sentences(text)
    return "".join(s[:m]).strip() if s else ""


def draft_context(mode, i, chunk_srcs, terms):
    """SOURCE-side context for the DRAFT/translator (§B3-бис wave 1): src-carryover + scene-state + lookahead.
    The THREE q3a modes SHARE the src-tail into the draft; they differ only in the editor's target-side."""
    blocks = []
    if mode.startswith("q3a") and i > 0:
        blocks.append(CTX_HEADER + f"\n[предыдущий фрагмент, источник]:\n{tail_sentences(chunk_srcs[i-1], K_CARRY)}")
        if mode == "q3a-state":
            st = scene_state_block(chunk_srcs[i], terms)
            if st:
                blocks.append(st)
    if mode == "q3b-look" and i < len(chunk_srcs) - 1:
        blocks.append(CTX_HEADER + f"\n[следующий фрагмент, источник]:\n{head_sentences(chunk_srcs[i+1], M_LOOK)}")
    return "\n\n".join(blocks)


def editor_context(mode, i, drafts, finals):
    """TARGET-side context for the EDITOR (§B3-бис wave 2): prev neighbor's draft (wave) or final
    (sequential) tail. q3a-src has NONE (source-carryover baseline); q3b-look targets the draft only."""
    if i == 0:
        return ""
    if mode == "q3a-draft" and drafts and drafts.get(i - 1):
        return CTX_HEADER + f"\n[предыдущий фрагмент, черновой перевод соседа]:\n{tail_sentences(drafts[i-1], K_CARRY)}"
    if mode in ("q3a-final", "q3a-state") and finals and finals.get(i - 1):
        return CTX_HEADER + f"\n[предыдущий фрагмент, итоговый перевод соседа]:\n{tail_sentences(finals[i-1], K_CARRY)}"
    return ""


def _reject(r):
    txt = (r.get("text") or "").strip()
    bad = bool(r.get("err")) or not txt or r.get("finish") not in ("stop", None)
    return txt, bad


def run_q3a(sp, bank, terms, chunk_srcs, base_drafts, base_finals, caps):
    """Q3a carrying block: ONE shared src-carryover DRAFT per chunk (§B3-бис: source-side -> wave-1 draft),
    then THREE editors differing ONLY in target-side neighbor context:
      q3a-src   : no target context      (source-carryover only)
      q3a-draft : flat-A0 prev-DRAFT tail (wave / draft-neighbor)
      q3a-final : flat-A0 prev-FINAL tail (sequential / final-neighbor)
    Sharing the draft removes the draft-noise confound between the 3 modes. Baseline = flat-A0-30 (backend).
    Resume-safe (skip non-empty existing files). Neighbors come from flat-A0-30 (parallel wave, no seq dep)."""
    tr = parse_template(BACKEND / "prompts/translator.md")
    ed = parse_template(BACKEND / "prompts/editor.md")
    tr_sys, ed_sys = fill(system_msg(tr)), fill(system_msg(ed))
    shared_dir = OUTDIR / "q3a-shared-draft"
    shared_dir.mkdir(parents=True, exist_ok=True)
    editor_modes = ["q3a-src", "q3a-draft", "q3a-final"]
    for m in editor_modes:
        (OUTDIR / m).mkdir(parents=True, exist_ok=True)
    active_ids = []
    limit = caps.get("max_chunks") or len(chunk_srcs)
    for i, src in enumerate(chunk_srcs):
        if i >= limit:
            break
        sticky = union_sticky(active_ids, len(active_ids))
        sel = bank.select(src, chapter=1, sticky_prev=sticky, budget_tokens=BUDGET_TOK)
        active_ids.append(sel.active_ids)
        gloss = render_glossary_block(sel.injected)
        edconstr = render_editor_constraint_block(sel.injected)
        dctx = draft_context("q3a-src", i, chunk_srcs, terms)   # src-carryover (SHARED across the 3 modes)

        # --- shared src-carryover draft (deepseek, temp 0.3), generated ONCE ---
        dpath = shared_dir / f"{i:02d}.draft.txt"
        if dpath.exists() and dpath.read_text(encoding="utf-8").strip():
            draft = dpath.read_text(encoding="utf-8")
        else:
            msgs = [{"role": "system", "content": tr_sys}]
            if gloss:
                msgs.append({"role": "system", "content": gloss})
            if dctx:
                msgs.append({"role": "system", "content": dctx})
            msgs.append({"role": "user", "content": fill(tr["user"], text=src)})
            ok, pred, why = sp.gate("deepseek-v4-flash", sum(len(m["content"]) for m in msgs) // 2, caps["draft_max"])
            if not ok:
                print(f"  [gate BLOCK] shared draft ch{i}: {why}"); return
            r = L.call("deepseek-v4-flash", msgs, temp=0.3, max_tokens=caps["draft_max"])
            sp.record({"mode": "q3a-shared", "stage": "draft", "chunk": i, "usage": r.get("usage"), "cost": r.get("cost", 0.0), "err": r.get("err"), "finish": r.get("finish")})
            draft, bad = _reject(r)
            if bad:
                print(f"  [BAD DRAFT] ch{i}: err={r.get('err')} finish={r.get('finish')}"); return
            dpath.write_text(draft, encoding="utf-8")

        # --- three editors on the shared draft (glm-5, temp 0.4), differing only in target context ---
        for mode in editor_modes:
            fpath = OUTDIR / mode / f"{i:02d}.final.txt"
            if fpath.exists() and fpath.read_text(encoding="utf-8").strip():
                continue
            ectx = editor_context(mode, i, base_drafts, base_finals)   # target-side (none/prev-draft/prev-final)
            emsgs = [{"role": "system", "content": ed_sys}]
            if edconstr:
                emsgs.append({"role": "system", "content": edconstr})
            if ectx:
                emsgs.append({"role": "system", "content": ectx})
            emsgs.append({"role": "user", "content": fill(ed["user"], text=src, draft=draft)})
            ok, pred, why = sp.gate("glm-5", sum(len(m["content"]) for m in emsgs) // 2, caps["edit_max"])
            if not ok:
                print(f"  [gate BLOCK] {mode} ch{i} edit: {why}"); return
            r = L.call("glm-5", emsgs, temp=0.4, max_tokens=caps["edit_max"])
            sp.record({"mode": mode, "stage": "edit", "chunk": i, "usage": r.get("usage"), "cost": r.get("cost", 0.0), "err": r.get("err"), "finish": r.get("finish")})
            final, bad = _reject(r)
            if bad:
                print(f"  [BAD FINAL] {mode} ch{i}: err={r.get('err')} finish={r.get('finish')}"); return
            fpath.write_text(final, encoding="utf-8")
            print(f"  ch{i} {mode}: final {len(final)}c  cost cum ${sp.total:.4f}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--modes", nargs="+", required=False,
                    help="a0 q3a-src q3a-draft q3a-final q3a-state q3b-look")
    ap.add_argument("--per-call-cap", type=float, default=0.05)
    ap.add_argument("--hard-cap", type=float, default=6.0)
    ap.add_argument("--max-chunks", type=int, default=0, help="sanity-gate: only first N chunks (0=all)")
    a = ap.parse_args()
    bank = build_bank(SEED)
    terms = seed_entities()
    chunk_srcs = greedy_chunks(S2.read_text(encoding="utf-8"))

    def load_base():
        """Carryover source + Q3 baseline = REAL backend flat-A0-30 (draft + editor final per chunk)."""
        import sqlite3
        con = sqlite3.connect(BASE_DB)
        d, f = {}, {}
        for ck, txt in con.execute("SELECT chunk_idx, response_text FROM checkpoints WHERE stage='draft' ORDER BY chunk_idx, attempt"):
            d[ck] = txt
        for ck, txt in con.execute("SELECT chunk_idx, response_text FROM checkpoints WHERE stage='edit' ORDER BY chunk_idx, attempt"):
            f[ck] = txt
        con.close()
        return d, f

    sp = L.Spender("/home/ubuntu/books/gu-zhenren/exp15/arm_costs.jsonl",
                   per_call_cap=a.per_call_cap, hard_cap=a.hard_cap)
    caps = {"draft_max": 16000, "edit_max": 12000, "max_chunks": a.max_chunks}
    base_d, base_f = load_base()
    n_needed = len(chunk_srcs)
    if len(base_d) < n_needed or len(base_f) < n_needed:
        print(f"  [ABORT] flat-A0-30 INCOMPLETE: draft {len(base_d)}/{n_needed}, final {len(base_f)}/{n_needed} "
              f"— the Q3 baseline + carryover source must be fully generated (regen must finish) first.")
        return
    lim = f", first {a.max_chunks} (sanity-gate)" if a.max_chunks else ""
    print(f"=== Q3a arms: shared src-carryover draft -> 3 editors (src/draft/final) over {len(chunk_srcs)} chunks{lim} ===")
    run_q3a(sp, bank, terms, chunk_srcs, base_d, base_f, caps)
    print(f"\narm spend this run: ${sp.total:.4f}")


if __name__ == "__main__":
    main()
