#!/usr/bin/env python3
"""從 wiki 的小說頁抽出所有「」對白，並用規則判定明確的說話者。

用法：extract_dialogues.py <wiki 目錄> <輸出 jsonl>

每個對白一筆：冊、章、前後文、說話者（規則無法判定則為 null）與規則依據。
規則只處理明確的句式；其餘交給 W4 的標註。輸出與原文放在 repo 之外。
"""
import json
import re
import sys
from collections import Counter
from pathlib import Path

ALIASES = json.loads((Path(__file__).with_name("characters.json")).read_text(encoding="utf-8"))
NAME_RE = "|".join(sorted(map(re.escape, ALIASES), key=len, reverse=True))
TOKEN_RE = re.compile(f"({NAME_RE})|我(?!們)")
QUOTE_RE = re.compile(r"「([^」]*)」")
SENT_END = "。！？!?"
CUE_RE = re.compile(r"說|問|叫|喊|答|道|聲音|低語|呢喃|嘀咕|開口|回|笑|哭|吼|嘟噥|喃喃")
# 「對我」「在我背後的」中的「我」是受詞，不是說話者。
OBJECT_RE = re.compile(
    r"(?:在我(?:的)?(?:背後|身後|旁邊|面前|耳邊|身邊)的?|(?:對著?|向|朝|跟|和|問|告訴|看著|瞪著|望著|回|給|讓|把|被)我(?:說|問)?)"
)
CJK_DIGITS = {"零": 0, "一": 1, "二": 2, "三": 3, "四": 4, "五": 5, "六": 6, "七": 7, "八": 8, "九": 9}
FILE_RE = re.compile(r"Sword Art Online (\d) .*? - 第(.+?)章(?: \((\d+)_(\d+)\))?\.md$")


def cn_number(s):
    if "十" in s:
        head, _, tail = s.partition("十")
        return (CJK_DIGITS[head] if head else 1) * 10 + (CJK_DIGITS[tail] if tail else 0)
    return CJK_DIGITS[s]


def strip_quotes(text):
    return QUOTE_RE.sub("", text)


def last_sentence(text):
    text = text.strip()
    parts = re.split(f"[{SENT_END}]", text)
    parts = [p for p in parts if p.strip()]
    return parts[-1].strip() if parts else ""


def first_sentence(text):
    return re.split(f"[{SENT_END}]", text.strip(), maxsplit=1)[0]


FRAME_RE = re.compile(r"^(?:面對|聽見|聽到|看見|看到|見到|對於|當|由於|因為|雖然|在|到了|一邊|就在)")


def pick_subject(clause):
    """回傳 (說話者別名, 全部候選, 是否高信心)。

    動詞前的主詞通常是句首第一個名字（跳過「面對…」「聽見…」之類的開頭子句）；
    只有一個候選且句中有說話動詞才算高信心，多個候選取第一個並降為中信心。
    """
    clause = OBJECT_RE.sub("", strip_quotes(clause))
    parts = [p for p in re.split("[，、,]", clause) if p.strip()]
    while len(parts) > 1 and FRAME_RE.match(parts[0]):
        parts = parts[1:]
    tokens = []
    for m in TOKEN_RE.finditer("，".join(parts)):
        alias = m.group(1) or "我"
        if alias not in tokens:
            tokens.append(alias)
    if not tokens:
        return None, [], False
    return tokens[0], tokens, len(tokens) == 1 and bool(CUE_RE.search(clause))


def resolve(alias):
    if alias == "我":
        return "kirito", "ingame"
    info = ALIASES[alias]
    return info["id"], info["mode"]


def attribute(before, after, prev_par):
    """回傳 (speaker, alias, confidence, rule, candidates)。"""
    b = before.strip()
    if b.endswith(("：", ":")):
        who, cands, cue = pick_subject(last_sentence(b))
        if who:
            return (*resolve(who)[:1], who, "high" if cue else "medium", "lead_in", cands)
        return None, None, None, None, cands
    if not b and prev_par and prev_par.rstrip().endswith(("：", ":")):
        who, cands, cue = pick_subject(last_sentence(prev_par))
        if who:
            return (*resolve(who)[:1], who, "high" if cue else "medium", "lead_in_prev", cands)
        return None, None, None, None, cands
    a = after.lstrip("，、 ")
    m = re.match(f"({NAME_RE}|我)", a)
    if m and not a.startswith("我們"):
        clause = first_sentence(a)
        if CUE_RE.search(clause):
            who, cands, _ = pick_subject(clause)
            if who == m.group(1):
                return (*resolve(who)[:1], who, "high", "trail", cands)
    return None, None, None, None, []


def trim(s, n=120):
    s = s.strip()
    return s if len(s) <= n else s[:n] + "…"


def load_chapters(wiki):
    pages = {}
    for f in sorted(Path(wiki).glob("*.md")):
        m = FILE_RE.search(f.name)
        if not m:
            print(f"略過：{f.name}", file=sys.stderr)
            continue
        vol, ch, part = int(m.group(1)), cn_number(m.group(2)), int(m.group(3) or 1)
        pages.setdefault((vol, ch), []).append((part, f))
    for key in sorted(pages):
        paras = []
        for part, f in sorted(pages[key]):
            paras += [(part, ln.strip()) for ln in f.read_text(encoding="utf-8").split("\n") if ln.strip()]
        yield key, paras


def extract(wiki):
    for (vol, ch), paras in load_chapters(wiki):
        seen = {}
        n = 0
        for pi, (part, par) in enumerate(paras):
            quotes = list(QUOTE_RE.finditer(par))
            prev_par = paras[pi - 1][1] if pi else ""
            next_par = paras[pi + 1][1] if pi + 1 < len(paras) else ""
            for qi, q in enumerate(quotes):
                lo = quotes[qi - 1].end() if qi else 0
                hi = quotes[qi + 1].start() if qi + 1 < len(quotes) else len(par)
                before, after = par[lo:q.start()], par[q.end():hi]
                speaker, alias, conf, rule, cands = attribute(before, after, prev_par if qi == 0 else "")
                text = q.group(1)
                n += 1
                dup = len(text) >= 8 and pi - seen.get(text, -99) <= 6
                seen[text] = pi
                rec = {
                    "id": f"v{vol}c{ch:02d}-{n:04d}", "volume": vol, "chapter": ch, "page": part, "para": pi,
                    "text": text, "before": trim(before), "after": trim(after),
                    "prev_par": trim(prev_par), "next_par": trim(next_par),
                    "speaker": speaker, "alias": alias, "confidence": conf, "rule": rule,
                    "mode": resolve(alias)[1] if alias else None,
                    "candidates": cands, "suspect_duplicate": dup,
                }
                yield rec


def main():
    wiki, out = sys.argv[1], Path(sys.argv[2])
    out.parent.mkdir(parents=True, exist_ok=True)
    stats, rules = Counter(), Counter()
    total = 0
    with out.open("w", encoding="utf-8") as fh:
        for rec in extract(wiki):
            fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
            total += 1
            stats[rec["speaker"] or "(未判定)"] += 1
            rules[rec["rule"] or "(無)"] += 1
    print(f"{total} 句 -> {out}")
    print("說話者：", dict(stats.most_common()))
    print("規則：", dict(rules.most_common()))


if __name__ == "__main__":
    main()