"""S3：从规范化 PDF 抽逐字坐标 → 切句 → 每句生成 quads（按行合并的矩形）和稳定编号，并做自动校验。

切句规则（原型级）：
  1. 同一行里，两个字之间的空隙大于 1.5 个字号时断开（表格单元格、制表位），得到「行片段」；
  2. 句子在 。！？； 处结束；行片段结束时，如果片段没写到右边距附近（段落末行）或者是表格单元格，也结束；
  3. 句子不跨页：跨页的句子拆成两段，各自有编号（出处本来就按页定位）。
编号：<文档>:p<页码>:s<页内序号>，规范化 PDF 不再重新生成，所以编号稳定。

校验：
  a. 几何回读：用 quads 裁剪原页，重新取字（字的中心落在矩形内），和句子原文比对；
  b. 文本保真：句子原文（去空白）能在 DOCX 原文件全文里找到，说明 PDF 抽出的文字没有乱码、错序。
用法：python -I sentences.py <规范化 PDF 目录> <DOCX 目录> <抽样数> <渲染张数>
"""
from __future__ import annotations

import json
import pathlib
import random
import re
import sys
import zipfile

import pdfplumber
import pypdfium2 as pdfium
from PIL import Image, ImageDraw

OUT = pathlib.Path(__file__).parent / "out"
END = set("。！？；")
WS = re.compile(r"[\s　]")


def line_fragments(page):
    """把页面上的字按行聚类，再按大空隙切成行片段。返回 [[char, ...], ...]，按阅读顺序。"""
    chars = [c for c in page.chars if c["text"].strip()]
    chars.sort(key=lambda c: (round(c["bottom"]), c["x0"]))
    lines, cur, base = [], [], None
    for c in chars:
        if base is not None and abs(c["bottom"] - base) > 0.45 * c["size"]:
            lines.append(cur); cur = []
        if not cur:
            base = c["bottom"]
        cur.append(c)
    if cur:
        lines.append(cur)
    frags = []
    for ln in lines:
        ln.sort(key=lambda c: c["x0"])
        f = [ln[0]]
        for a, b in zip(ln, ln[1:]):
            if b["x0"] - a["x1"] > 1.5 * a["size"]:
                frags.append(f); f = []
            f.append(b)
        frags.append(f)
    return frags


def split_sentences(doc_id, page, right_edge):
    sents, cur = [], []
    frags = line_fragments(page)
    for fi, frag in enumerate(frags):
        for c in frag:
            cur.append((fi, c))
            if c["text"] in END:
                sents.append(cur); cur = []
        short = frag[-1]["x1"] < right_edge - 2 * frag[-1]["size"]
        if cur and short:
            sents.append(cur); cur = []
    if cur:
        sents.append(cur)
    out = []
    for i, s in enumerate(sents, 1):
        quads = {}
        for fi, c in s:
            q = quads.setdefault(fi, [c["x0"], c["top"], c["x1"], c["bottom"]])
            q[0], q[1] = min(q[0], c["x0"]), min(q[1], c["top"])
            q[2], q[3] = max(q[2], c["x1"]), max(q[3], c["bottom"])
        text = "".join(c["text"] for _, c in s)
        out.append({"id": f"{doc_id}:p{page.page_number}:s{i}", "page": page.page_number, "text": text,
                    "quads": [[round(v, 2) for v in q] for q in quads.values()]})
    return out


def reread(page, quads):
    """几何回读：字的中心落在任一 quad 里就算在内，按行、按 x 排序拼接。"""
    got = []
    for c in page.chars:
        if not c["text"].strip():
            continue
        cx, cy = (c["x0"] + c["x1"]) / 2, (c["top"] + c["bottom"]) / 2
        for qi, (x0, t, x1, b) in enumerate(quads):
            if x0 <= cx <= x1 and t <= cy <= b:
                got.append((qi, c["x0"], c["text"])); break
    got.sort()
    return "".join(t for _, _, t in got)


def docx_text(path):
    """DOCX 正文文字。WPS 生成的旧版 DOC 转 DOCX 后，部分文字是 <w:sym w:char="XXXX"/>（符号字符），也要算进来。"""
    x = zipfile.ZipFile(path).read("word/document.xml").decode()
    parts = re.findall(r'<w:t[^>]*>([^<]*)</w:t>|<w:sym [^>]*w:char="([0-9A-Fa-f]+)"', x)
    return WS.sub("", "".join(t or chr(int(h, 16)) for t, h in parts))


NUMBERING = re.compile(r"^(\d+[.、．]|[（(]\d+[)）]|\d+\))")   # Word 自动编号由渲染生成，不在 DOCX 文字里
PAGE_FIELD = re.compile(r"^第\d+页(共\d+页)?$")                  # 页眉页脚的页码域


def main():
    pdf_dir, docx_dir, n_sample, n_render = pathlib.Path(sys.argv[1]), pathlib.Path(sys.argv[2]), int(sys.argv[3]), int(sys.argv[4])
    all_sents, stats = [], {}
    for pdf_path in sorted(pdf_dir.glob("S*.pdf")):
        sid = pdf_path.stem
        with pdfplumber.open(pdf_path) as pdf:
            sents = []
            for page in pdf.pages:
                right = max((c["x1"] for c in page.chars), default=page.width)
                sents += split_sentences(sid, page, right)
        (OUT / f"{sid}-sentences.json").write_text(json.dumps(sents, ensure_ascii=False))
        stats[sid] = {"sentences": len(sents), "multi_line": sum(len(s["quads"]) > 1 for s in sents)}
        all_sents += sents
    rng = random.Random(int(sys.argv[5]) if len(sys.argv) > 5 else 20261103)
    # 只抽 ≥ 8 个字的句子：太短的是页码、编号，不会作为出处
    pool = [s for s in all_sents if len(WS.sub("", s["text"])) >= 8]
    # 页码域是页眉页脚，不会作为出处；只统计、不抽样
    page_fields = sum(bool(PAGE_FIELD.match(WS.sub("", s["text"]))) for s in pool)
    pool = [s for s in pool if not PAGE_FIELD.match(WS.sub("", s["text"]))]
    sample = rng.sample(pool, n_sample)
    docx_cache, results = {}, []
    for s in sample:
        sid, page_no = s["id"].split(":")[0], s["page"]
        with pdfplumber.open(pdf_dir / f"{sid}.pdf") as pdf:
            back = reread(pdf.pages[page_no - 1], s["quads"])
        if sid not in docx_cache:
            docx_cache[sid] = docx_text(docx_dir / f"{sid}.docx")
        norm = WS.sub("", s["text"])
        body = NUMBERING.sub("", norm)
        results.append({"id": s["id"], "text": s["text"][:60], "lines": len(s["quads"]),
                        "geom_ok": WS.sub("", back) == norm, "back": back[:60] if WS.sub("", back) != norm else "",
                        "in_docx": norm in docx_cache[sid],
                        "in_docx_ignoring_autonumber": body in docx_cache[sid]})
    geom = sum(r["geom_ok"] for r in results); fid = sum(r["in_docx"] for r in results)
    summary = {"docs": stats, "total_sentences": len(all_sents), "sample": n_sample,
               "geom_ok": geom, "text_in_docx": fid,
               "text_in_docx_ignoring_autonumber": sum(r["in_docx_ignoring_autonumber"] for r in results),
               "page_field_sentences_excluded": page_fields,
               "multi_line_in_sample": sum(r["lines"] > 1 for r in results)}
    print(json.dumps(summary, ensure_ascii=False))
    for r in results:
        if not (r["geom_ok"] and r["in_docx_ignoring_autonumber"]):
            print("  FAIL", json.dumps(r, ensure_ascii=False))
    (OUT / "validation.json").write_text(json.dumps({"summary": summary, "results": results}, ensure_ascii=False, indent=1))
    # 渲染高亮图：优先挑跨行的句子，配色用设计规范的蓝铅笔浅底 + 蓝铅笔描边
    picks = [s for s in sample if len(s["quads"]) > 1][: n_render // 2] + [s for s in sample if len(s["quads"]) == 1][: n_render - n_render // 2]
    for s in picks:
        sid, page_no = s["id"].split(":")[0], s["page"]
        scale = 2.0
        img = pdfium.PdfDocument(str(pdf_dir / f"{sid}.pdf"))[page_no - 1].render(scale=scale).to_pil().convert("RGBA")
        overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
        d = ImageDraw.Draw(overlay)
        for x0, t, x1, b in s["quads"]:
            box = [x0 * scale - 2, t * scale - 2, x1 * scale + 2, b * scale + 2]
            d.rectangle(box, fill=(231, 237, 251, 150), outline=(47, 91, 211, 255), width=2)
        Image.alpha_composite(img, overlay).convert("RGB").save(OUT / f"highlight-{s['id'].replace(':', '_')}.png")


if __name__ == "__main__":
    main()