#!/usr/bin/env python3
"""Terse-dialogue precision of the coverage gate (POLYGON package-2 Task 5; closes the untested
dialogue-dense case of the webnovel slice).

WHY: exp07 measured the coverage gate at 0 FP/57 on CLASSIC prose. The webnovel slice was supposed
to close the DIALOGUE-DENSE case (terse quoted replies are the noisiest input for sent_cov/len_ratio)
but did NOT — its dialogue detector counted only lines that START with a quote, while webnovels attribute
dialogue INLINE (方源道：“…”), so dialogue-dense chunks read as prose (<0.5) and never formed a terse
bucket. This harness fixes detection + segmentation and measures the gate's FALSE excision_suspect rate
on GOOD (judged-complete) translations of a ≥20-chunk terse-dialogue bucket from 蛊真人. That rate gates
removing the 1-flag tolerance (D12/Q4) and the acceptance thresholds (exp02/07) for D18.

METHOD:
  1. Pull dialogue-dense 节 sections from the book; segment RESPECTING source line breaks (each source
     line = a paragraph unit — never glue replies into one blob), group consecutive lines into ~target
     chunks. dialogue_density = fraction of lines carrying a QUOTED SPAN (“…”/「…」/『…』), not just a
     leading quote. terse = short average quoted-span length.
  2. Keep chunks with density >= --min-density → the terse bucket (aim >= 20).
  3. Translate each with a reliable, complete translator (default grok non-reasoning: no echo mine).
  4. classify_output → excision_suspect (sent_cov<0.75 or len_ratio<corridor-low).
  5. Cross-family completeness judge (D13.3) on EVERY chunk → ground truth complete? So:
       false_excision = gate flagged excision_suspect BUT judge says complete (no omission).
  6. Report: excision rate on terse bucket, false-excision rate, and the split vs prose control.

Full provenance: raw translations + judge JSON + usage + UTC under data/dialogue_precision/raw_<stamp>/.
The book text stays OUT of git (Р8); only counts go to the journal/exp.

Usage:
  eval/.venv/bin/python eval/dialogue_precision.py --dry-run          # select bucket, no API
  eval/.venv/bin/python eval/dialogue_precision.py                    # full run
"""
from __future__ import annotations
import argparse, json, re, sys, time
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))
import refusal_bench as rb

BOOK = Path("/home/ubuntu/books/gu-zhenren/guzhenren-utf8.txt")
OUTDIR = ROOT / "data" / "dialogue_precision"
SEC_RE = re.compile(r"^第[0-9一二三四五六七八九十百千零两]{1,10}节")
QUOTE_SPAN = re.compile(r"[“「『][^”」』]*[”」』]")   # a quoted span anywhere in the line (inline attribution ok)
PROVIDERS = {p["name"]: p for p in json.loads((ROOT / "providers.json").read_text())["providers"]}
FAMILY = {"deepseek": "deepseek", "grok": "xai", "gemini": "google", "openai": "openai",
          "glm": "zhipu", "kimi": "moonshot", "mistral": "mistral"}

TRANSLATE_SYSTEM = ("Ты профессиональный литературный переводчик. Переведи фрагмент китайского "
                    "произведения на русский язык. Сохрани ВСЕ реплики и детали без пропусков и "
                    "смягчений; стиль — живой литературный русский. Выведи ТОЛЬКО перевод.")

COMPLETE_SYSTEM = (
    "Ты — строгий контролёр ПОЛНОТЫ художественного перевода (не стиля). Тебе дают ИСХОДНЫЙ фрагмент "
    "(китайский) и его ПЕРЕВОД на русский. Проверь, не ПРОПУЩЕНА ли какая-либо реплика, предложение "
    "или значимая деталь (omission). Верни СТРОГО JSON: "
    '{"complete": <true|false>, "n_omissions": <int>, "omitted": [<кратко что пропущено>], '
    '"notes": "<кратко>"}. complete=true — ничего существенного не пропущено (мелкие стилевые сжатия '
    "допустимы); complete=false — есть реальный пропуск реплики/предложения/детали.")


def load_sections() -> list[dict]:
    lines = BOOK.read_text(encoding="utf-8").split("\n")
    marks = [(i, l.strip()) for i, l in enumerate(lines) if SEC_RE.match(l.strip())]
    secs = []
    for k, (i, title) in enumerate(marks):
        end = marks[k + 1][0] if k + 1 < len(marks) else len(lines)
        body = [l.strip() for l in lines[i + 1:end] if l.strip()]  # paragraph units = source lines
        secs.append({"n": k + 1, "title": title, "lines": body})
    return secs


def dialogue_density(chunk_lines: list[str]) -> float:
    if not chunk_lines:
        return 0.0
    d = sum(1 for l in chunk_lines if QUOTE_SPAN.search(l))
    return round(d / len(chunk_lines), 3)


def terse_score(chunk_lines: list[str]) -> float:
    """Mean quoted-span length over the chunk's quoted spans (lower = terser)."""
    spans = [m.group(0) for l in chunk_lines for m in QUOTE_SPAN.finditer(l)]
    return round(sum(len(s) for s in spans) / max(1, len(spans)), 1) if spans else 0.0


def segment_lines(lines: list[str], target: int) -> list[list[str]]:
    """Group consecutive source lines into ~target-char chunks WITHOUT gluing across the line grain."""
    chunks, buf, tot = [], [], 0
    for l in lines:
        if buf and tot + len(l) > target:
            chunks.append(buf); buf, tot = [], 0
        buf.append(l); tot += len(l)
    if buf:
        chunks.append(buf)
    return chunks


def build_buckets(target: int, min_density: float, want: int) -> tuple[list[dict], list[dict]]:
    """Scan the book, segment, split chunks into terse-dialogue (density>=min) and prose control."""
    terse, prose = [], []
    for s in load_sections():
        for ci, cl in enumerate(segment_lines(s["lines"], target)):
            text = "\n".join(cl)
            if len(text) < 400:
                continue
            dens = dialogue_density(cl)
            rec = {"id": f"sec{s['n']}#{ci}", "sec": s["n"], "n_lines": len(cl), "chars": len(text),
                   "dialogue_density": dens, "terse_len": terse_score(cl), "text": text}
            (terse if dens >= min_density else prose).append(rec)
    terse.sort(key=lambda r: (-r["dialogue_density"], r["terse_len"]))
    prose.sort(key=lambda r: r["dialogue_density"])
    return terse[:max(want, 20)], prose[:max(want // 2, 10)]


def retry_call(p, system, user, timeout=200, attempts=3):
    last = (None, "unknown", {})
    for i in range(attempts):
        t, err, usage = rb.call_provider(p, system, user, timeout=timeout)
        if not err or not re.match(r"(http\|(503|429|500|502|504))|transport\|", err or ""):
            return t, err, usage
        last = (t, err, usage); time.sleep(2 * (i + 1))
    return last


def judge_complete(source, translation, judge_name):
    p = PROVIDERS[judge_name]
    user = (f"ИСХОДНЫЙ ФРАГМЕНТ (китайский):\n{source}\n\nПЕРЕВОД НА РУССКИЙ:\n{translation}\n\nВерни JSON.")
    txt, err, usage = retry_call({"name": judge_name, **p}, COMPLETE_SYSTEM, user, timeout=200)
    rec = {"judge_model": p["model"], "usage": usage or {},
           "ts": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")}
    if err or not txt:
        rec["error"] = err or "empty"; return rec
    m = re.search(r"\{.*\}", txt, re.S)
    if not m:
        rec["error"], rec["raw"] = "no-json", txt[:300]; return rec
    try:
        rec.update(json.loads(m.group(0)))
    except json.JSONDecodeError:
        rec["error"], rec["raw"] = "bad-json", txt[:300]
    return rec


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--translator", default="grok", help="complete translator (no echo mine)")
    ap.add_argument("--judge", default="gemini", help="cross-family completeness judge (D13.3)")
    ap.add_argument("--target", type=int, default=900)
    ap.add_argument("--min-density", type=float, default=0.5)
    ap.add_argument("--want", type=int, default=24)
    ap.add_argument("--max-prose", type=int, default=10, help="prose control chunks to also run")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    if FAMILY.get(args.judge) == FAMILY.get(args.translator):
        sys.exit(f"D13.3: judge {args.judge} shares family with translator {args.translator}")
    terse, prose = build_buckets(args.target, args.min_density, args.want)
    prose = prose[:args.max_prose]
    print(f"terse bucket: {len(terse)} chunks (density>={args.min_density}); prose control: {len(prose)}", file=sys.stderr)
    print(f"  terse density range {terse[-1]['dialogue_density']}..{terse[0]['dialogue_density']} "
          f"mean terse_len {round(sum(r['terse_len'] for r in terse)/max(1,len(terse)),1)}", file=sys.stderr)
    if args.dry_run:
        for r in terse[:24]:
            print(f"  {r['id']} dens={r['dialogue_density']} terse_len={r['terse_len']} lines={r['n_lines']} chars={r['chars']}", file=sys.stderr)
        return

    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    raw_dir = OUTDIR / f"raw_{stamp}"; raw_dir.mkdir(parents=True, exist_ok=True)
    tp = PROVIDERS[args.translator]

    def run_bucket(bucket, label):
        rows = []
        for r in bucket:
            t0 = time.time()
            out, err, usage = retry_call({"name": args.translator, **tp}, TRANSLATE_SYSTEM, r["text"], timeout=240)
            out = out or ""
            verdict = rb.classify_output(r["text"], out, err, rb.EXPECT_LEN_RATIO.get("zh", (2.2, 4.2)))
            jr = judge_complete(r["text"], out, args.judge)
            complete = jr.get("complete")
            flagged = verdict["verdict"] == "excision_suspect"
            false_excision = bool(flagged and complete is True)
            (raw_dir / f"{label}_{r['id'].replace('#','_')}.txt").write_text(
                f"# translator={tp['model']} id={r['id']} density={r['dialogue_density']} verdict={verdict['verdict']}\n"
                f"# judge_complete={complete} judge={jr.get('judge_model')} usage={usage}\n"
                f"# --- SRC ---\n{r['text']}\n\n# --- TRANSLATION ---\n{out}\n", encoding="utf-8")
            row = {**{k: r[k] for k in ("id", "sec", "dialogue_density", "terse_len", "n_lines", "chars")},
                   "verdict": verdict["verdict"], "sent_cov": verdict.get("sent_cov"),
                   "len_ratio": verdict.get("len_ratio"), "flagged_excision": flagged,
                   "judge_complete": complete, "judge_omissions": jr.get("n_omissions"),
                   "false_excision": false_excision, "judge": jr, "translation": out,
                   "err": err, "usage": usage, "elapsed_s": round(time.time() - t0, 1)}
            rows.append(row)
            print(f"  [{label}] {r['id']} dens={r['dialogue_density']} verdict={verdict['verdict']} "
                  f"complete={complete} sent_cov={verdict.get('sent_cov')} len_ratio={verdict.get('len_ratio')} "
                  f"{'FALSE_EXCISION' if false_excision else ''}", file=sys.stderr)
        return rows

    terse_rows = run_bucket(terse, "terse")
    prose_rows = run_bucket(prose, "prose")

    def summarize(rows):
        n = len(rows)
        flagged = sum(1 for r in rows if r["flagged_excision"])
        complete = sum(1 for r in rows if r["judge_complete"] is True)
        false_exc = sum(1 for r in rows if r["false_excision"])
        judged = sum(1 for r in rows if isinstance(r["judge_complete"], bool))
        return {"n": n, "flagged_excision": flagged, "judge_complete": complete, "judged": judged,
                "false_excision": false_exc,
                "false_excision_rate_of_complete": round(false_exc / max(1, complete), 3),
                "excision_rate": round(flagged / max(1, n), 3)}

    ts, ps = summarize(terse_rows), summarize(prose_rows)
    meta = {"run_utc": stamp, "translator": tp["model"], "judge": PROVIDERS[args.judge]["model"],
            "target": args.target, "min_density": args.min_density,
            "terse_summary": ts, "prose_summary": ps}
    out_path = OUTDIR / f"dialogue_precision_{stamp}.json"
    out_path.write_text(json.dumps({"meta": meta, "terse": terse_rows, "prose": prose_rows}, ensure_ascii=False, indent=2))

    print(f"\n=== TERSE-DIALOGUE COVERAGE-GATE PRECISION ({ts['n']} chunks) ===", file=sys.stderr)
    print(f"  terse: excision_flagged={ts['flagged_excision']}/{ts['n']} "
          f"judged_complete={ts['judge_complete']}/{ts['judged']} "
          f"FALSE_excision={ts['false_excision']} → false-excision-rate-of-complete={ts['false_excision_rate_of_complete']}", file=sys.stderr)
    print(f"  prose control: excision_flagged={ps['flagged_excision']}/{ps['n']} "
          f"FALSE_excision={ps['false_excision']} rate={ps['false_excision_rate_of_complete']}", file=sys.stderr)
    print(f"  [exp07 classic baseline: 0 FP/57]", file=sys.stderr)
    print(f"\nwrote {out_path}\nraw: {raw_dir}", file=sys.stderr)


if __name__ == "__main__":
    main()