#!/usr/bin/env python3
"""exp14b — армы батареи смысл-трапов + P1a↔F-disc. ПРОМПТ константен (дискурс), варьируем
editor-модель: C=glm-5 / F=gpt-5.4 / X=grok-4.3(reasoning-ON). Общий flash-черновик.

Реюз exp14-выходов где есть (не пере-оплата). Per-call кап ОБЯЗАТЕЛЕН. ЧЕСТНЫЙ учёт ошибок
(finish=length/empty/http/retry — раздельно; НЕ «0 ошибок»). $ персист.

Использование: exp14b_arms.py [--dry]
"""
from __future__ import annotations
import argparse, json, shutil, 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

RERUN = Path("/home/ubuntu/books/gu-zhenren/rerun")
RECS = {(r["chapter"], r["chunk_idx"]): r for r in json.load(open(RERUN / "records.json"))}
INJ = json.load(open(Path("/home/ubuntu/books/gu-zhenren/exp14") / "injection_blocks.json"))
EX14 = Path("/home/ubuntu/books/gu-zhenren/exp14/arms")
OUT = Path("/home/ubuntu/books/gu-zhenren/exp14b"); OUT.mkdir(parents=True, exist_ok=True)
ARMS = OUT / "arms"; ARMS.mkdir(exist_ok=True)

# базовые (пре-рег §1) + добавленные семьи-редакторы (расширение панели floor-теста, прозрачно)
MODEL = {"C": "glm-5", "F": "gpt-5.4", "X": "grok-4.3",
         "D": "deepseek-v4-pro", "K": "kimi-k2.6", "M": "mistral-large-2512"}
EX14_ARM = {"C": "P1a", "F": "F-disc", "X": "X-disc"}  # соответствие для реюза (только базовые)
CAPS = {"glm-5": 0.08, "gpt-5.4": 0.40, "grok-4.3": 0.40,
        "deepseek-v4-pro": 0.06, "kimi-k2.6": 0.20, "mistral-large-2512": 0.10}

# battery chunks + stage2 (для head-to-head C vs F)
BATTERY = ["7.0", "8.1", "11.1", "17.0", "19.0", "10.1", "6.0", "16.0", "2.1", "9.1"]
STAGE2 = ["17.0", "17.1", "18.0", "18.1", "18.2", "19.0", "19.1"]
# что каждому арму нужно покрыть
NEED = {
    "C": sorted(set(BATTERY) | set(STAGE2)),
    "F": sorted(set(BATTERY) | set(STAGE2)),   # F на stage2 уже есть в exp14 → реюз
    "X": sorted(set(BATTERY)),
    "D": sorted(set(BATTERY)),   # deepseek-v4-pro — только батарея (floor-тест)
    "K": sorted(set(BATTERY)),   # kimi-k2.6
    "M": sorted(set(BATTERY)),   # mistral-large-2512
}


def esys():
    return P.editor_system("discourse")


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--dry", action="store_true"); a = ap.parse_args()
    sp = C.Spender(OUT / "costs.jsonl", CAPS)
    errs = {"length": 0, "empty": 0, "http": 0, "content_filter": 0, "other": 0, "retried_ok": 0}
    sysd = esys()
    for arm, chunks in NEED.items():
        d = ARMS / arm; d.mkdir(exist_ok=True)
        model = MODEL[arm]
        for cid in chunks:
            outp = d / f"{cid}.txt"
            if outp.exists():
                continue
            # РЕЮЗ exp14 (только базовые C/F/X — D/K/M новых семей в exp14 не было)
            if arm in EX14_ARM:
                ex = EX14 / EX14_ARM[arm] / f"{cid}.txt"
                if ex.exists():
                    shutil.copy(ex, outp)
                    print(f"  [reuse exp14] {arm} {cid} ← {EX14_ARM[arm]}")
                    continue
            ch, ck = cid.split("."); ch = int(ch); ck = int(ck)
            r = RECS.get((ch, ck))
            if not r:
                print(f"  [MISS rec] {cid}"); continue
            einj = INJ.get(f"{ch}/{ck}", {}).get("editor_constraint", "")
            user = P.editor_user(r["source"], r["draft"])
            msgs = C.messages_with_injection(sysd, einj, user)
            fam = "deepseek" if model.startswith("deepseek") else "glm"
            in_est = C.count_tokens(sysd + einj + user, fam)
            s = C.spec(model, temp=0.4)
            ok, pred = sp.guard(model, in_est, s["max_tokens"])
            print(f"  {arm} {cid}: in≈{in_est} predict=${pred:.4f} cap=${CAPS[model]} {'OK' if ok else 'OVER'}")
            if a.dry or not ok:
                continue
            # до 2 ретраев на transient (honest error accounting)
            text, err, usage, attempts = "", None, {}, 0
            for attempt in range(3):
                attempts += 1
                text, err, usage = C.call(s, msgs, timeout=360)
                sp.record({"arm": arm, "model": model, "keyenv": s["keyenv"], "id": cid,
                           "attempt": attempts, "err": err, "usage": usage})
                if not err:
                    if attempt > 0: errs["retried_ok"] += 1
                    break
                # классификация ошибки
                if "length" in (err or ""): errs["length"] += 1
                elif "empty" in (err or ""): errs["empty"] += 1
                elif "http" in (err or ""): errs["http"] += 1
                elif "content_filter" in (err or ""): errs["content_filter"] += 1
                else: errs["other"] += 1
                if "http" in (err or "") or "transport" in (err or ""):
                    time.sleep(3); continue  # transient → retry
                break  # length/empty/filter → не ретраим слепо
            if err:
                print(f"    ERR {err[:60]} (attempts={attempts})"); continue
            outp.write_text(text, encoding="utf-8")
            print(f"    ok {len(text)}ch")
    print(f"\n=== ЧЕСТНЫЙ учёт ошибок: {errs} ===")
    print(f"=== spend by key: {dict((k,round(v,4)) for k,v in sp.by_key.items())} ===")


if __name__ == "__main__":
    main()