#!/usr/bin/env python3
"""exp14 — раннер editor-армов (нарезка × промпт + аблации). Прямые вызовы, ручная упаковка.

Per-call predicted-cost кап (НЕ group-level). Персист usage/cost в exp14/costs.jsonl.
P0 = reuse records.json final ($0). Черновик держим общим (flash-draft из records) — варьируем
editor-стадию (модель × промпт × нарезка × reference).

Использование: exp14_arms.py --arm P1a [--chapters trap|stage2] [--model glm-5] [--dry]
Армы: P0 P1a P1b P1c A-layout A-2pass A-ref-prev A-ref-next A-ref-both  (+ frontier: F-base F-disc)
"""
from __future__ import annotations
import argparse, json, re, sys, time
from pathlib import Path

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

MAT = json.load(open(C.DIAG / "material.json"))
ARMSDIR = C.DIAG / "arms"; ARMSDIR.mkdir(exist_ok=True)
CAPS = {"glm-5": 0.08, "glm-5.1": 0.10, "deepseek-v4-flash": 0.03, "deepseek-v4-pro": 0.06,
        "gpt-5.4": 0.40, "gemini-3.1-pro-preview": 0.30, "grok-4.3": 0.40,
        "gpt-5-mini": 0.20, "kimi-k2.6": 0.20, "mistral-large-2512": 0.10}
SP = C.Spender(C.DIAG / "costs.jsonl", CAPS)


def deformat_draft(draft: str) -> str:
    """A-layout: снять построчную/пустострочную разбивку черновика (нейтральные сепараторы)."""
    lines = [l.strip() for l in draft.split("\n") if l.strip()]
    return " ".join(lines)  # один поток; редактор не якорится на форме черновика


def chunk_units(scope: str):
    """Возвращает список единиц {id, source, draft, final_p0, editor_constraint, chapter, chunk_idx}."""
    if scope == "trap":
        return [dict(id=f"{u['chapter']}.{u['chunk_idx']}", **u) for u in MAT["trap_chunks"]]
    if scope == "stage2":
        out = []
        for ch, chunks in MAT["stage2"].items():
            for u in chunks:
                out.append(dict(id=f"{ch}.{u['chunk_idx']}", chapter=int(ch), **u))
        return out
    raise ValueError(scope)


def chapter_units(scope: str):
    """Whole-chapter единицы: клеим source/draft чанков главы + chapter-level инъекция."""
    inj = json.load(open(C.DIAG / "injection_blocks.json"))
    units = chunk_units(scope)
    by_ch = {}
    for u in units:
        by_ch.setdefault(u["chapter"], []).append(u)
    out = []
    for ch, us in by_ch.items():
        if len(us) < 2:  # whole-chapter тест осмыслен только когда все чанки главы в наборе (ch8/11/24)
            continue
        us = sorted(us, key=lambda x: x["chunk_idx"])
        src = "\n\n".join(x["source"] for x in us)
        drf = "\n\n".join(x["draft"] for x in us)
        eb = inj.get(f"{ch}/ALL", {}).get("editor_constraint", "")
        out.append(dict(id=f"{ch}.ALL", chapter=ch, chunk_idx="ALL",
                        source=src, draft=drf, editor_constraint=eb))
    return out


def run_editor(arm: str, variant: str, units, model: str, reference_by=None,
               two_pass=False, deformat=False, dry=False):
    outdir = ARMSDIR / arm; outdir.mkdir(exist_ok=True)
    sys_disc = P.editor_system(variant)
    total_pred = 0.0
    for u in units:
        outp = outdir / f"{u['id']}.txt"
        if outp.exists():
            print(f"  [skip exists] {arm} {u['id']}"); continue
        draft = deformat_draft(u["draft"]) if deformat else u["draft"]
        reference = reference_by.get(u["id"], "") if reference_by else ""
        user = P.editor_user(u["source"], draft, reference)
        msgs = C.messages_with_injection(sys_disc, u.get("editor_constraint", ""), user)
        in_est = C.count_tokens(sys_disc + u.get("editor_constraint", "") + user,
                                "deepseek" if model.startswith("deepseek") else "glm")
        s = C.spec(model, temp=0.4)
        ok, pred = SP.guard(model, in_est, s["max_tokens"])
        total_pred += pred
        print(f"  {arm} {u['id']}: in≈{in_est}tok predict=${pred:.4f} cap=${CAPS[model]} {'OK' if ok else 'OVER-CAP'}")
        if dry:
            continue
        if not ok:
            print(f"    !! OVER per-call cap — skipped"); continue
        t0 = time.time()
        text, err, usage = C.call(s, msgs, timeout=360)
        dt = round(time.time() - t0, 1)
        rec = {"arm": arm, "model": model, "keyenv": s["keyenv"], "id": u["id"],
               "variant": variant, "err": err, "latency_s": dt, "usage": usage}
        c = SP.record(rec)
        if err:
            print(f"    ERR {err[:80]} (${c:.4f}, {dt}s)")
            continue
        if two_pass:  # second narrow pass: target-only literary reflow of the just-edited RU
            sys2 = P.editor_system("discourse")  # target-only reflow uses discourse core
            u2 = ("Ниже — русский перевод. Перевёрстай его в естественные русские абзацы по "
                  "правилам дискурс-перевёрстки; НЕ меняй смысл, НЕ добавляй и НЕ удаляй атомы, "
                  "НЕ сверяйся с иностранным исходником (его нет):\n\n" + text)
            msgs2 = [{"role": "system", "content": sys2}, {"role": "user", "content": u2}]
            in2 = C.count_tokens(sys2 + u2, "glm")
            ok2, pred2 = SP.guard(model, in2, s["max_tokens"])
            if ok2:
                text2, err2, usage2 = C.call(s, msgs2, timeout=360)
                SP.record({"arm": arm, "model": model, "keyenv": s["keyenv"], "id": u["id"],
                           "variant": "reflow-pass2", "err": err2, "usage": usage2})
                if not err2:
                    text = text2
        outp.write_text(text, encoding="utf-8")
        print(f"    ok {len(text)}chars ${c:.4f} {dt}s → {outp.name}")
    print(f"  [{arm}] predicted total ≈ ${total_pred:.4f}; spent-by-key: {SP.by_key}")


def run_p0(scope):
    outdir = ARMSDIR / "P0"; outdir.mkdir(exist_ok=True)
    for u in chunk_units(scope):
        (outdir / f"{u['id']}.txt").write_text(u["final_p0"], encoding="utf-8")
    print(f"[P0] reused {len(chunk_units(scope))} finals ($0)")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", required=True)
    ap.add_argument("--scope", default="trap", choices=["trap", "stage2"])
    ap.add_argument("--model", default="glm-5")
    ap.add_argument("--dry", action="store_true")
    a = ap.parse_args()

    if a.arm == "P0":
        run_p0(a.scope); return
    if a.arm == "P1a":
        run_editor("P1a", "discourse", chunk_units(a.scope), a.model, dry=a.dry)
    elif a.arm == "P1b":
        run_editor("P1b", "baseline", chapter_units(a.scope), a.model, dry=a.dry)
    elif a.arm == "P1c":
        run_editor("P1c", "discourse", chapter_units(a.scope), a.model, dry=a.dry)
    elif a.arm == "A-layout":
        run_editor("A-layout", "discourse", chunk_units(a.scope), a.model, deformat=True, dry=a.dry)
    elif a.arm == "A-2pass":
        run_editor("A-2pass", "baseline", chunk_units(a.scope), a.model, two_pass=True, dry=a.dry)
    elif a.arm in ("A-ref-prev", "A-ref-next", "A-ref-both"):
        ref = build_reference(a.arm)
        run_editor(a.arm, "discourse", [u for u in chunk_units(a.scope) if u["id"] in ref],
                   a.model, reference_by=ref, dry=a.dry)
    elif a.arm in ("F-base", "F-disc"):
        variant = "baseline" if a.arm == "F-base" else "discourse"
        run_editor(a.arm, variant, chunk_units(a.scope), a.model, dry=a.dry)
    else:
        raise SystemExit(f"unknown arm {a.arm}")


def build_reference(kind: str) -> dict:
    """±1 reference для кросс-граничных трапов (T5 24.1←24.0, T6 11.1←11.0). prev/next/both."""
    units = {u["id"]: u for u in chunk_units("trap")}
    pairs = [("24.0", "24.1"), ("11.0", "11.1")]  # (prev_chunk, target_chunk)
    ref = {}
    for prev, tgt in pairs:
        parts = []
        if kind in ("A-ref-prev", "A-ref-both"):
            ptail = "\n".join(units[prev]["source"].split("\n")[-6:])
            ttail = "\n".join(units[prev]["final_p0"].split("\n")[-6:])
            parts.append("[ПРЕД-ИСХОДНИК]\n" + ptail + "\n[ПРЕД-ПЕРЕВОД]\n" + ttail)
        if kind in ("A-ref-next", "A-ref-both"):
            # next-source: следующий чанк того же trap-набора, если есть (для 24.1/11.1 нет — пропуск)
            pass
        if parts:
            ref[tgt] = "\n\n".join(parts)
    return ref


if __name__ == "__main__":
    main()