#!/usr/bin/env python3
"""蛊真人 (Gu Zhenren) 18+ corpus builder helper — section index + violence-density scoring +
fragment extraction, for the D14.4 explicit-benchmark corpus (exp10, Task 3a).

The book text and all derivatives are OUTSIDE git (Р8: private processing legal). This helper reads
the UTF-8 working copy (iconv'd from the GB18030 original) and:
  - splits it into 节 (sections) by the '第…节' markers;
  - scores each section by a violence/dark lexicon density (metadata only — NEVER dumps bodies);
  - extracts a paragraph-bounded fragment (~target chars) from a chosen section;
  - writes JSONL records in the refusal_corpus schema.

The category is ALWAYS violence/dark here (level 2). Level-3 (minors in a sexual context) is out of
scope by construction and MUST be screened by the human/agent selecting fragments — this book is
violence/scheming-driven (女主：无, no female lead), so erotica coverage is expected to be SPARSE;
report that honestly rather than forcing weak fragments.

Usage:
  eval/.venv/bin/python eval/gu_corpus_build.py index                  # section count + volumes
  eval/.venv/bin/python eval/gu_corpus_build.py candidates --top 40    # top violence-dense sections (metadata)
  eval/.venv/bin/python eval/gu_corpus_build.py show --sec 1 --chars 1200   # print ONE fragment (for review)
  eval/.venv/bin/python eval/gu_corpus_build.py extract --spec spec.json --out <file>.jsonl
"""
from __future__ import annotations
import argparse, json, re, sys
from pathlib import Path

BOOK = Path("/home/ubuntu/books/gu-zhenren/guzhenren-utf8.txt")
SEC_RE = re.compile(r"^第[0-9一二三四五六七八九十百千零两]{1,10}[节卷]")

# Violence / dark / gore lexicon (Simplified). Density = weighted hits / section chars.
VIO_LEX = {
    "杀": 3, "血": 2, "尸": 3, "尸体": 4, "鲜血": 3, "死": 1, "惨": 2, "惨叫": 4, "割": 3,
    "断": 1, "炼": 1, "毒": 1, "痛": 1, "痛苦": 2, "骨": 1, "撕": 3, "咬": 2, "爆": 1,
    "碎": 1, "残忍": 4, "残酷": 3, "虐": 4, "折磨": 4, "血肉": 4, "剖": 4, "剐": 5,
    "肢解": 6, "屠": 4, "刑": 2, "尖叫": 3, "颅": 3, "食人": 6, "生吞": 5, "剥皮": 6,
}


def load_sections() -> list[dict]:
    if not BOOK.exists():
        sys.exit(f"missing UTF-8 working copy: {BOOK}\n  iconv -f GB18030 -t UTF-8 "
                 f"/home/ubuntu/books/gu-zhenren/guzhenren-gb18030.txt > {BOOK}")
    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 = []
    # keep only 节 (section) markers as content units; 卷 (volume) headers just tag the volume
    node_marks = [(i, t) for (i, t) in marks if "节" in t]
    for k, (i, title) in enumerate(node_marks):
        end = node_marks[k + 1][0] if k + 1 < len(node_marks) else len(lines)
        body = "\n".join(lines[i + 1:end]).strip()
        secs.append({"n": k + 1, "line": i, "title": title, "body": body, "chars": len(body)})
    return secs


def score(body: str) -> tuple[float, dict]:
    hits = {}
    total = 0
    for kw, w in VIO_LEX.items():
        c = body.count(kw)
        if c:
            hits[kw] = c
            total += c * w
    density = total / max(1, len(body))
    return density, hits


def para_fragment(body: str, target: int) -> str:
    paras = [p.strip() for p in re.split(r"\n\s*\n|\n", body) if p.strip()]
    buf, tot = [], 0
    for p in paras:
        if buf and tot + len(p) > target:
            break
        buf.append(p)
        tot += len(p)
    return "\n".join(buf) if buf else body[:target]


def cmd_index(_):
    secs = load_sections()
    print(f"sections(节): {len(secs)}; total chars: {sum(s['chars'] for s in secs):,}")
    print(f"median section chars: {sorted(s['chars'] for s in secs)[len(secs)//2]}")


def cmd_candidates(args):
    secs = load_sections()
    scored = []
    for s in secs:
        if s["chars"] < 600:
            continue
        d, hits = score(s["body"])
        top = sorted(hits.items(), key=lambda kv: -kv[1] * VIO_LEX[kv[0]])[:6]
        scored.append((d, s["n"], s["title"], s["chars"], top))
    scored.sort(key=lambda x: -x[0])
    print(f"# top {args.top} violence-dense sections (metadata only)")
    print("# density | sec# | chars | top-kw(count) | title")
    for d, n, title, chars, top in scored[:args.top]:
        kw = " ".join(f"{k}×{c}" for k, c in top)
        print(f"{d:.4f} | {n:5} | {chars:5} | {kw:28} | {title}")


def cmd_show(args):
    secs = load_sections()
    s = next((x for x in secs if x["n"] == args.sec), None)
    if not s:
        sys.exit(f"no section {args.sec}")
    frag = para_fragment(s["body"], args.chars)
    print(f"# sec {s['n']} '{s['title']}' ({len(frag)} chars of {s['chars']})")
    print(frag)


def cmd_extract(args):
    """spec.json: [{"id","sec","chars","category","note"}...] → JSONL refusal_corpus records."""
    secs = {x["n"]: x for x in load_sections()}
    spec = json.loads(Path(args.spec).read_text(encoding="utf-8"))
    out = []
    for item in spec:
        s = secs[item["sec"]]
        frag = para_fragment(s["body"], item.get("chars", 1200))
        out.append({"id": item["id"], "lang": "zh", "level": 2,
                    "category": item.get("category", "violence"),
                    "source": f"蛊真人 节{item['sec']} «{s['title']}» (владелец, вне git)",
                    "license": "владелец — приватная обработка (Р8), вне git",
                    "expected": "translate",
                    "note": item.get("note", ""), "text": frag})
    Path(args.out).write_text("\n".join(json.dumps(r, ensure_ascii=False) for r in out) + "\n",
                              encoding="utf-8")
    print(f"wrote {len(out)} records → {args.out}")
    for r in out:
        print(f"  {r['id']} sec-based {len(r['text'])} chars [{r['category']}]")


def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    sub.add_parser("index").set_defaults(fn=cmd_index)
    c = sub.add_parser("candidates"); c.add_argument("--top", type=int, default=40); c.set_defaults(fn=cmd_candidates)
    sh = sub.add_parser("show"); sh.add_argument("--sec", type=int, required=True); sh.add_argument("--chars", type=int, default=1200); sh.set_defaults(fn=cmd_show)
    ex = sub.add_parser("extract"); ex.add_argument("--spec", required=True); ex.add_argument("--out", required=True); ex.set_defaults(fn=cmd_extract)
    args = ap.parse_args()
    args.fn(args)


if __name__ == "__main__":
    main()