#!/usr/bin/env python3
"""exp15 — extract A0/A0' from the real backend DBs+logs and compute:
  - per-chunk SOURCE (recovered from draft user-msg heads), DRAFT, FINAL (editor) texts
  - greedy cut offsets (A0 boundaries — the Q1 common-denominator boundaries) + est_out distribution
  - deterministic KPI per whole-anchor (mean_spp, 1-sent-para frac, dialogue-dash frac, CJK-leak,
    glossary/term consistency vs seed)
  - A0<->A0' deterministic divergence (chunk count, boundary offsets, KPI deltas, term consistency,
    char-level similarity) = PRELIMINARY noise floor (the judged trap floor is measured with the judges)

Chunk SOURCE recovery reuses validate_injection.recover_chunk_sources (draft user-msg head -> offset in
S2'.txt). FINAL text = the edit-stage response_text per chunk (the authoritative editor output).
"""
from __future__ import annotations
import json
import re
import sqlite3
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from validate_injection import load_exchanges, recover_chunk_sources
from chunker import est_out, est_tokens
import regex

S2 = Path("/home/ubuntu/books/gu-zhenren/exp15/S2prime.txt")
_RE_HAN = regex.compile(r"\p{Han}")
_RE_CYR = regex.compile(r"[А-Яа-яЁё]")


def final_texts(db):
    """chunk_idx -> editor (final) response_text, latest attempt. S2' is one chapter."""
    con = sqlite3.connect(db)
    rows = con.execute(
        "SELECT chunk_idx, attempt, response_text FROM checkpoints WHERE stage='edit' ORDER BY chunk_idx, attempt"
    ).fetchall()
    con.close()
    out = {}
    for ck, att, txt in rows:
        out[ck] = txt   # later attempt overwrites
    return out


def draft_texts(db):
    con = sqlite3.connect(db)
    rows = con.execute(
        "SELECT chunk_idx, attempt, response_text FROM checkpoints WHERE stage='draft' ORDER BY chunk_idx, attempt"
    ).fetchall()
    con.close()
    out = {}
    for ck, att, txt in rows:
        out[ck] = txt
    return out


def load_anchor(db, log):
    src_text = S2.read_text(encoding="utf-8")
    exchanges = load_exchanges(Path(log))
    srcs = recover_chunk_sources(exchanges, src_text, whole_file=False)  # {(chap,ck): source}
    finals = final_texts(db)
    drafts = draft_texts(db)
    chunks = []
    for (chap, ck) in sorted(srcs.keys()):
        src = srcs[(chap, ck)]
        chunks.append({"chapter": chap, "chunk": ck, "source": src,
                       "draft": drafts.get(ck), "final": finals.get(ck),
                       "src_est_out": est_out(src) if src else None,
                       "src_est_tokens": est_tokens(src) if src else None})
    return chunks, src_text


def cut_offsets(chunks):
    """Rune offsets in S2' where each greedy cut falls (chunk starts, excluding 0)."""
    offs, cur = [], 0
    for c in chunks:
        if c["source"] is None:
            return None
        if cur > 0:
            offs.append(cur)
        cur += len(c["source"])
    return offs


# --- deterministic KPI ------------------------------------------------------------------------------
def paragraphs(text):
    return [p for p in re.split(r"\n\s*\n|\n", text) if p.strip()]


def sentences_ru(text):
    return [s for s in re.split(r"(?<=[.!?…])\s+", text) if s.strip()]


def kpi(final_text):
    paras = paragraphs(final_text)
    spp = [len(sentences_ru(p)) for p in paras] or [0]
    mean_spp = sum(spp) / len(spp)
    one_sent = sum(1 for x in spp if x == 1) / len(spp) if spp else 0
    dash_paras = sum(1 for p in paras if p.lstrip().startswith(("—", "–", "-")))
    dash_frac = dash_paras / len(paras) if paras else 0
    # CJK leak: Han runes in the RU output NOT immediately after a Cyrillic word in parens gloss
    # (owner decision #3: «Кулак Ло Ханя (羅漢拳)» balanced Han-in-parens is NOT a leak). Crude: count
    # Han runes that are inside parentheses -> whitelisted; Han outside parens -> leak.
    han_total = len(_RE_HAN.findall(final_text))
    han_in_parens = 0
    for m in re.finditer(r"[（(]([^）)]*)[）)]", final_text):
        han_in_parens += len(_RE_HAN.findall(m.group(1)))
    han_leak = han_total - han_in_parens
    return {"n_paras": len(paras), "mean_spp": round(mean_spp, 2), "one_sent_para_frac": round(one_sent, 3),
            "dialogue_dash_frac": round(dash_frac, 3), "han_total": han_total,
            "han_in_parens": han_in_parens, "han_leak_outside_parens": han_leak,
            "ru_chars": len(_RE_CYR.findall(final_text))}


def term_consistency(final_text, seed_path):
    """LTCR-style: for each approved seed term present in S2', is its canonical dst (or a decl form)
    present in the final? Reports coverage and any missing (a drift/omission signal, deterministic)."""
    import yaml
    d = yaml.safe_load(Path(seed_path).read_text(encoding="utf-8"))
    src_text = S2.read_text(encoding="utf-8")
    hits, miss = [], []
    for t in d["terms"]:
        src = t.get("src"); dst = (t.get("dst") or "").strip()
        if not src or not dst or (t.get("status") or "approved") != "approved":
            continue
        variants = [src] + [a.get("alias") for a in (t.get("aliases") or []) if a.get("alias")]
        if not any(v in src_text for v in variants):
            continue
        forms = set([dst] + list((t.get("decl") or {}).get("forms") or []))
        # match the distinctive first content token of the dst (case-insensitive)
        present = any(f.lower() in final_text.lower() for f in forms)
        if not present:
            # fallback: distinctive stem (first word >=4 chars)
            toks = [w for w in re.findall(r"[А-Яа-яЁё]+", dst) if len(w) >= 4]
            present = any(w.lower()[:5] in final_text.lower() for w in toks)
        (hits if present else miss).append(dst)
    return {"n_terms_in_source": len(hits) + len(miss), "present": len(hits),
            "missing": miss, "coverage": round(len(hits) / (len(hits) + len(miss)), 3) if (hits or miss) else None}


def main():
    seed = "/home/ubuntu/books/gu-zhenren/guzhenren-seed-v2.yaml"
    AH = "/home/ubuntu/books/gu-zhenren/exp15/anchors"
    out = {}
    for tag in ("flat_a0", "flat_a0prime"):
        chunks, _ = load_anchor(f"{AH}/{tag}.db", f"{AH}/{tag}.log.json")
        whole_final = "\n\n".join(c["final"] for c in chunks if c["final"])
        offs = cut_offsets(chunks)
        src_outs = [c["src_est_out"] for c in chunks if c["src_est_out"]]
        out[tag] = {
            "n_chunks": len(chunks),
            "cut_offsets": offs,
            "src_est_out": {"mean": round(sum(src_outs) / len(src_outs), 1), "min": round(min(src_outs), 1),
                            "max": round(max(src_outs), 1), "per_chunk": [round(x, 1) for x in src_outs]},
            "kpi": kpi(whole_final),
            "term_consistency": term_consistency(whole_final, seed),
            "final_chars": len(whole_final),
        }
        Path(f"{AH}/{tag}.final.txt").write_text(whole_final, encoding="utf-8")
        Path(f"{AH}/{tag}.chunks.json").write_text(
            json.dumps([{k: c[k] for k in ("chunk", "source", "draft", "final", "src_est_out")} for c in chunks],
                       ensure_ascii=False, indent=2), encoding="utf-8")
    # A0<->A0' deterministic floor
    a, b = out["flat_a0"], out["flat_a0prime"]
    floor = {
        "same_chunk_count": a["n_chunks"] == b["n_chunks"],
        "same_cut_offsets": a["cut_offsets"] == b["cut_offsets"],
        "n_chunks": [a["n_chunks"], b["n_chunks"]],
        "kpi_delta": {k: round(a["kpi"][k] - b["kpi"][k], 3) for k in ("mean_spp", "one_sent_para_frac",
                      "dialogue_dash_frac", "han_leak_outside_parens") if isinstance(a["kpi"][k], (int, float))},
        "term_coverage": [a["term_consistency"]["coverage"], b["term_consistency"]["coverage"]],
        "final_chars": [a["final_chars"], b["final_chars"]],
    }
    result = {"anchors": out, "deterministic_floor_A0_vs_A0prime": floor}
    Path(f"{AH}/anchors_analysis.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps({"a0_chunks": a["n_chunks"], "a0prime_chunks": b["n_chunks"],
                      "a0_cut_offsets": a["cut_offsets"], "same_boundaries": floor["same_cut_offsets"],
                      "a0_src_est_out_mean": a["src_est_out"]["mean"],
                      "floor": floor}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()
