#!/usr/bin/env python3
"""exp12 — извлечение per-chunk (source/draft/final) из приёмочного store + выбор 3 чистых глав.

$0. Никаких вызовов. Источники (ВНЕ git):
  - store копия:  /home/ubuntu/books/gu-zhenren/diag/store-copy.db (draft=deepseek, edit=glm-5)
  - параллель:    /home/ubuntu/books/gu-zhenren/acceptance/guzhenren-25ch-parallel.txt ([ОРИГ]/[ПЕРЕВОД])

Выход:
  - diag/chunks.json — per (chapter,chunk): source, draft(A0), final(A1), dispositions
  - diag/selection.json — метрики выбора глав + 3 выбранные главы (пре-регистрированное правило)

Пре-регистрированное правило выбора (frozen, см. 12-quality-diagnosis.md §Пре-регистрация):
  frame = ЧИСТЫЕ главы (все чанки disposition=ok на draft И edit, flag_reason пуст).
  метрика на ИСХОДНИКЕ zh: D = доля символов внутри кавычек “…” (U+201C…U+201D) от не-пробельных.
  выбор: high=argmax D, low=argmin D, mid=ближайшая к медиане D. Тай-брейк — меньший номер главы.
"""
from __future__ import annotations
import json, re, sqlite3, sys
from pathlib import Path

DIAG = Path("/home/ubuntu/books/gu-zhenren/diag")
DB = DIAG / "store-copy.db"
PARALLEL = Path("/home/ubuntu/books/gu-zhenren/acceptance/guzhenren-25ch-parallel.txt")

CH_HDR = re.compile(r"【\s*Глава\s+(\d+)\s*】")
CHUNK_HDR = re.compile(r"^---\s*чанк\s+(\d+)\s*\[(.+?)\]\s*---")
LQUOTE, RQUOTE = "“", "”"  # “ ”


def parse_parallel(path: Path) -> dict:
    """→ {(ch,chunk): {'source': str, 'final_parallel': str, 'flag': str}}"""
    out: dict = {}
    ch = None
    chunk = None
    mode = None  # 'orig' | 'per'
    buf_src: list[str] = []
    buf_per: list[str] = []

    def flush():
        if ch is not None and chunk is not None:
            out[(ch, chunk)] = {
                "source": "\n".join(buf_src).strip(),
                "final_parallel": "\n".join(buf_per).strip(),
                "flag": flag,
            }

    flag = ""
    for line in path.read_text(encoding="utf-8").splitlines():
        m = CH_HDR.search(line)
        if m:
            flush(); ch = int(m.group(1)); chunk = None; mode = None
            buf_src, buf_per = [], []
            continue
        m = CHUNK_HDR.match(line.strip())
        if m:
            flush()
            chunk = int(m.group(1)); flag = m.group(2).strip()
            mode = None; buf_src, buf_per = [], []
            continue
        s = line.strip()
        if s == "[ОРИГ]":
            mode = "orig"; continue
        if s == "[ПЕРЕВОД]":
            mode = "per"; continue
        if s.startswith("[⚠") or s.startswith("=====") or not s and mode is None:
            continue
        if mode == "orig":
            buf_src.append(line)
        elif mode == "per":
            buf_per.append(line)
    flush()
    return out


def load_store(db: Path) -> dict:
    """→ {(ch,chunk): {'draft':..,'edit':..,'draft_disp':..,'edit_disp':..,'draft_flag':..,'edit_flag':..}}"""
    con = sqlite3.connect(db)
    con.row_factory = sqlite3.Row
    rows = con.execute("""
      SELECT cs.chapter, cs.chunk_idx, cs.stage, cs.disposition, cs.flag_reason,
             cp.model_actual, cp.response_text
      FROM chunk_status cs
      LEFT JOIN checkpoints cp ON cp.request_hash = cs.final_hash
      ORDER BY cs.chapter, cs.chunk_idx, cs.stage
    """).fetchall()
    con.close()
    out: dict = {}
    for r in rows:
        key = (r["chapter"], r["chunk_idx"])
        d = out.setdefault(key, {})
        st = r["stage"]
        d[st] = r["response_text"] or ""
        d[f"{st}_disp"] = r["disposition"]
        d[f"{st}_flag"] = r["flag_reason"]
        d[f"{st}_model"] = r["model_actual"]
    return out


def dialogue_density(src: str) -> float:
    nonspace = re.sub(r"\s", "", src)
    if not nonspace:
        return 0.0
    inside = 0
    depth = 0
    for c in src:
        if c == LQUOTE:
            depth += 1; continue
        if c == RQUOTE:
            if depth > 0:
                depth -= 1
            continue
        if depth > 0 and not c.isspace():
            inside += 1
    return inside / len(nonspace)


def main() -> None:
    par = parse_parallel(PARALLEL)
    store = load_store(DB)

    # merge per chunk
    chunks = {}
    keys = sorted(set(par) | set(store))
    for k in keys:
        ch, ci = k
        p = par.get(k, {})
        s = store.get(k, {})
        chunks[f"{ch}/{ci}"] = {
            "chapter": ch, "chunk": ci,
            "source": p.get("source", ""),
            "draft": s.get("draft", ""),       # A0
            "final": s.get("edit", ""),        # A1 (glm-5 mono)
            "final_parallel": p.get("final_parallel", ""),
            "draft_disp": s.get("draft_disp", "?"),
            "edit_disp": s.get("edit_disp", "?"),
            "draft_flag": s.get("draft_flag", ""),
            "edit_flag": s.get("edit_flag", ""),
        }

    # sanity: store edit vs parallel final match
    mism = [k for k, v in chunks.items()
            if v["final"] and v["final_parallel"]
            and v["final"].strip()[:120] != v["final_parallel"].strip()[:120]]

    # per-chapter clean classification + dialogue density
    chap = {}
    for k, v in chunks.items():
        c = v["chapter"]
        d = chap.setdefault(c, {"chunks": 0, "clean": True, "src": [], "flags": []})
        d["chunks"] += 1
        d["src"].append(v["source"])
        if v["draft_disp"] != "ok" or v["edit_disp"] != "ok":
            d["clean"] = False
            d["flags"].append(f"{v['draft_flag']}|{v['edit_flag']}")

    metrics = {}
    for c, d in chap.items():
        full_src = "\n".join(d["src"])
        metrics[c] = {
            "clean": d["clean"],
            "chunks": d["chunks"],
            "src_chars": len(re.sub(r"\s", "", full_src)),
            "dialogue_density": round(dialogue_density(full_src), 4),
        }

    clean = {c: m for c, m in metrics.items() if m["clean"]}
    ds = sorted((m["dialogue_density"], c) for c, m in clean.items())
    dens = [x[0] for x in ds]
    median = dens[len(dens) // 2] if len(dens) % 2 else (dens[len(dens)//2 - 1] + dens[len(dens)//2]) / 2

    high = max(clean, key=lambda c: (clean[c]["dialogue_density"], -c))
    low = min(clean, key=lambda c: (clean[c]["dialogue_density"], c))
    mid = min((c for c in clean if c not in (high, low)),
              key=lambda c: (abs(clean[c]["dialogue_density"] - median), c))

    selection = {
        "rule": "clean-frame; D=share of chars in “…”; high=argmax, low=argmin, mid=nearest-median; tie=min chapter#",
        "clean_chapters": sorted(clean),
        "median_D": round(median, 4),
        "selected": {
            "high_dialogue": {"chapter": high, **clean[high]},
            "mid": {"chapter": mid, **clean[mid]},
            "low_exposition": {"chapter": low, **clean[low]},
        },
        "all_clean_metrics": {str(c): clean[c] for c in sorted(clean)},
    }

    DIAG.mkdir(exist_ok=True)
    (DIAG / "chunks.json").write_text(json.dumps(chunks, ensure_ascii=False, indent=1), encoding="utf-8")
    (DIAG / "selection.json").write_text(json.dumps(selection, ensure_ascii=False, indent=1), encoding="utf-8")

    print(f"chunks extracted: {len(chunks)}  | store-parallel final mismatches (first120): {len(mism)}")
    if mism:
        print("  mism keys:", mism[:10])
    print(f"clean chapters ({len(clean)}): {sorted(clean)}")
    print(f"median D = {median:.4f}")
    print("\nDialogue density (clean, sorted):")
    for dval, c in ds:
        star = " <-HIGH" if c == high else " <-LOW" if c == low else " <-MID" if c == mid else ""
        print(f"  ch{c:>2}: D={dval:.4f}  chunks={clean[c]['chunks']}  src_chars={clean[c]['src_chars']}{star}")
    print(f"\nSELECTED: high=ch{high} (D={clean[high]['dialogue_density']}), "
          f"mid=ch{mid} (D={clean[mid]['dialogue_density']}), low=ch{low} (D={clean[low]['dialogue_density']})")


if __name__ == "__main__":
    main()