#!/usr/bin/env python3
"""exp14 §кривая — размер-чанка ↔ качество ↔ биллинг ↔ ретраи.

Пере-нарезка главы при разных бюджетах (жадная упаковка целых абзацев-строк до char-порога)
+ whole-chapter. Пайплайн draft(flash thinking-ON)→edit(glm-5 дискурс). Снимаем РЕАЛЬНЫЕ
in/out/reasoning-токены, латентность, ошибки-по-причинам, output-token-cost, KPI-абзацы.
Chapter-level инъекция едина для всех сегментов (изолируем РАЗМЕР, не глоссарий).

Бюджет чанка отчитывается в ВЫХОДНЫХ токенах (реальный токенайзер glm-4.6). Оптимизация —
retry-adjusted output-cost. Оговорка: у этой вебновеллы плоская структура (1 предл./строка,
мало сцен-маркеров) → политика границ «сцена vs жадная» слабо различима; меряем жадную + whole.

Использование: exp14_sizecurve.py [--chapters 17,13] [--dry]
"""
from __future__ import annotations
import argparse, json, re, statistics, 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(C.DIAG / "injection_blocks.json"))
SIZES = [800, 1600, 3200, 99999]  # zh-chars; 99999 = whole-chapter
OUT = C.DIAG / "sizecurve"; OUT.mkdir(exist_ok=True)
CAPS = {"deepseek-v4-flash": 0.03, "glm-5": 0.08}


def chapter_source(ch):
    cks = sorted(c for (cc, c) in RECS if cc == ch)
    return "\n".join(RECS[(ch, c)]["source"] for c in cks)


def rechunk(src, budget):
    """Жадная упаковка целых строк (абзацев) до budget zh-символов; строку не рвём."""
    lines = [l for l in src.split("\n") if l.strip()]
    segs, cur, n = [], [], 0
    for l in lines:
        if cur and n + len(l) > budget:
            segs.append("\n".join(cur)); cur, n = [], 0
        cur.append(l); n += len(l)
    if cur:
        segs.append("\n".join(cur))
    return segs


def paras(t): return [p for p in re.split(r"\n\s*\n", t.strip()) if p.strip()]
def sents(p): return [s for s in re.split(r"(?<=[.!?…])\s+", p.strip()) if s.strip()]
def kpi(t):
    narr = [p for p in paras(t) if not p.strip().startswith("—")]
    sp = [len(sents(p)) for p in narr] or [0]
    return round(statistics.mean(sp), 2), round(sum(1 for x in sp if x == 1)/len(sp), 3)


def run():
    ap = argparse.ArgumentParser(); ap.add_argument("--chapters", default="17,13")
    ap.add_argument("--dry", action="store_true"); a = ap.parse_args()
    chapters = [int(x) for x in a.chapters.split(",")]
    sp = C.Spender(C.DIAG / "sizecurve_costs.jsonl", CAPS)
    tsys = P.translator_system(); esys = P.editor_system("discourse")
    rows = []
    for ch in chapters:
        src = chapter_source(ch)
        einj = INJ.get(f"{ch}/ALL", {}).get("editor_constraint", "")
        ginj = INJ.get(f"{ch}/ALL", {}).get("bilingual_glossary", "")
        for budget in SIZES:
            segs = rechunk(src, budget)
            label = "whole" if budget == 99999 else f"{budget}c"
            tag = f"{ch}-{label}"
            outp = OUT / f"{tag}.txt"
            if outp.exists():
                print(f"  [skip] {tag}"); continue
            print(f"\n=== ch{ch} size={label}: {len(segs)} сегментов ===")
            if a.dry:
                for s in segs:
                    print(f"    seg {len(s)}zh-chars")
                continue
            final_parts, in_tot, out_tot, reason_tot, seg_costs, errs, lat = [], 0, 0, 0, 0.0, [], 0.0
            for i, seg in enumerate(segs):
                # DRAFT (flash, thinking ON)
                ds = C.spec("deepseek-v4-flash", temp=0.3, max_tokens=8000)
                dmsgs = C.messages_with_injection(tsys, ginj, P.translator_user(seg))
                t0 = time.time()
                draft, derr, dusage = C.call(ds, dmsgs, timeout=300)
                sp.record({"model": "deepseek-v4-flash", "keyenv": ds["keyenv"], "tag": tag,
                           "seg": i, "stage": "draft", "err": derr, "usage": dusage})
                if derr:
                    errs.append(f"draft:{derr[:30]}"); draft = seg  # фолбэк — исходник (эхо-гард поймает)
                # EDIT (glm-5 discourse)
                es = C.spec("glm-5", temp=0.4, max_tokens=8000)
                emsgs = C.messages_with_injection(esys, einj, P.editor_user(seg, draft))
                final, eerr, eusage = C.call(es, emsgs, timeout=360)
                c = sp.record({"model": "glm-5", "keyenv": es["keyenv"], "tag": tag, "seg": i,
                               "stage": "edit", "err": eerr, "usage": eusage})
                lat += round(time.time() - t0, 1)
                if eerr:
                    errs.append(f"edit:{eerr[:30]}"); final = draft
                final_parts.append(final)
                in_tot += (dusage.get("prompt_tokens", 0) or 0) + (eusage.get("prompt_tokens", 0) or 0)
                out_tot += (dusage.get("completion_tokens", 0) or 0) + (eusage.get("completion_tokens", 0) or 0)
            full = "\n\n".join(final_parts)
            outp.write_text(full, encoding="utf-8")
            ru_out_tok = C.count_tokens(full, "glm")
            mean_spp, share1 = kpi(full)
            han = sum(1 for ch_ in full if "一" <= ch_ <= "鿿")
            row = dict(chapter=ch, size=label, n_seg=len(segs), in_tok=in_tot, out_tok=out_tot,
                       ru_out_tok=ru_out_tok, cost=round(sp.by_key.get("ZAI_API_KEY", 0)+sp.by_key.get("DEEPSEEK_API_KEY", 0), 4),
                       mean_spp=mean_spp, share_1sent=share1, han=han, errs=errs, latency_s=round(lat, 1))
            rows.append(row)
            print(f"    segs={len(segs)} in={in_tot} out={out_tot} ru_out_tok={ru_out_tok} "
                  f"mean_spp={mean_spp} share1={share1} CJK={han} errs={len(errs)} lat={lat:.0f}s")
    (OUT / "sizecurve.json").write_text(json.dumps(rows, ensure_ascii=False, indent=1), encoding="utf-8")
    if rows:
        print("\n=== КРИВАЯ (per chapter×size) ===")
        print(f"{'ch-size':<12}{'segs':>5}{'out_tok':>9}{'mean_spp':>9}{'share1':>8}{'CJK':>5}{'errs':>5}")
        for r in rows:
            print(f"{r['chapter']}-{r['size']:<8}{r['n_seg']:>5}{r['out_tok']:>9}{r['mean_spp']:>9}{r['share_1sent']:>8}{r['han']:>5}{len(r['errs']):>5}")


if __name__ == "__main__":
    run()