Community_Nursing_Atlas/pipeline/toc.py

"""목차 페이지에서 장·절·항 트리를 추출."""
import hashlib
import re

import pymupdf

from .pdfutil import page_lines

LEADER = re.compile(r"(?:j{3,}|[·․\.…]{3,}|·\s?·)[\sj·․\.…]*")
TRAIL_PAGE = re.compile(r"(?:_|\s|^)(\d{1,3})$")
ROMAN = "ⅠⅡⅢⅣⅤⅥⅦⅧⅨⅩ"

LEVELS = {"chapter": 1, "part": 2, "section": 3, "item": 4, "subitem": 5}


def _split_page(text: str):
    t = LEADER.sub(" ", text).strip()
    m = TRAIL_PAGE.search(t)
    if m and (m.start() > 0 or t.startswith("_")):
        return t[: m.start()].strip(" _"), int(m.group(1))
    return t, None


def _classify(style: str, text: str, size: float):
    """(kind, number, title) 또는 None."""
    if style == "cn2026":
        for kind, pat in (
            ("chapter", r"^(\d+)장\s*(.+)$"),
            ("section", r"^(\d+)절\s*(.+)$"),
            ("item", r"^(\d+)\.\s*(.+)$"),
            ("subitem", r"^(\d+)\)\s*(.+)$"),
        ):
            m = re.match(pat, text)
            if m:
                return kind, m.group(1), m.group(2).strip()
        return None
    if style == "pacific":
        m = re.match(r"^제(\d+)장\s*(.*)$", text)
        if m:
            return "chapter", m.group(1), m.group(2).strip()
        m = re.match(rf"^([{ROMAN}])\.?\s*(.*)$", text)
        if m:
            return "part", str(ROMAN.index(m.group(1)) + 1), m.group(2).strip()
        for kind, pat in (
            ("section", r"^제(\d+)절\s*(.+)$"),
            ("item", r"^(\d+)\.\s*(.+)$"),
            ("subitem", r"^(\d+)\)\s*(.+)$"),
        ):
            m = re.match(pat, text)
            if m:
                return kind, m.group(1), m.group(2).strip()
        return None
    if style == "healthedu":
        m = re.match(r"^(\d+)\.\s*(.+)$", text)
        if m and size >= 8.7:
            return "chapter", m.group(1), re.sub(r"(?<=[가-힣]) (?=[가-힣]$)", "", m.group(2).strip())
        m = re.match(r"^\((\d+)\)\s*(.+)$", text)
        if m:
            return "subitem", m.group(1), m.group(2).strip()
        m = re.match(r"^(\d+)\s+(.+)$", text)
        if m and size < 8.2:
            return "section", m.group(1), m.group(2).strip()
        return None
    raise ValueError(style)


def extract_toc(doc: pymupdf.Document, book: dict) -> list[dict]:
    style = book["toc_style"]
    nodes: list[dict] = []
    stack: list[dict] = []
    pending_title_for = None  # 제목이 다음 줄에 오는 장/부
    seq = 0
    for pi in book["toc_pages"]:
        for ln in page_lines(doc, pi):
            if ln.size >= 16 or ln.y < 60:  # 쪽 제목·머리말
                continue
            text, page = _split_page(ln.text)
            if not text:
                continue
            if text in ("Contents", "목차", "w"):
                continue
            c = _classify(style, text, ln.size)
            if c is None:
                if pending_title_for is not None and not pending_title_for["title"]:
                    pending_title_for["title"] = text
                    if page is not None and pending_title_for["page"] is None:
                        pending_title_for["page"] = page
                    pending_title_for = None
                elif nodes and nodes[-1]["page"] is None:
                    nodes[-1]["title"] = (nodes[-1]["title"] + " " + text).strip()
                    if page is not None:
                        nodes[-1]["page"] = page
                continue
            kind, num, title = c
            level = LEVELS[kind]
            while stack and stack[-1]["level"] >= level:
                stack.pop()
            seq += 1
            node = {
                # 공개 화면 주소에 교재가 드러나지 않도록 해시 ID
                "id": "s_" + hashlib.sha1(f"{book['book_id']}|{seq}".encode()).hexdigest()[:10],
                "book_id": book["book_id"],
                "kind": kind,
                "level": level,
                "number": num,
                "title": title,
                "page": page,
                "parent_id": stack[-1]["id"] if stack else None,
            }
            # 같은 장이 목차 여러 쪽에 걸쳐 반복 표기된 경우 건너뜀
            same = [n for n in nodes if n["kind"] == "chapter" and n["number"] == num
                    and (not title or not n["title"] or n["title"] == title)]
            if kind == "chapter" and same:
                seq -= 1
                stack = [same[0]]
                continue
            nodes.append(node)
            stack.append(node)
            pending_title_for = node if not title else None
    _fill_pages(nodes)
    return nodes


def _fill_pages(nodes: list[dict]) -> None:
    """페이지가 없는 노드는 첫 자식의 페이지를 쓰고, 끝 페이지를 계산."""
    by_parent: dict = {}
    for n in nodes:
        by_parent.setdefault(n["parent_id"], []).append(n)
    for n in reversed(nodes):
        if n["page"] is None:
            kids = [k for k in by_parent.get(n["id"], []) if k["page"] is not None]
            if kids:
                n["page"] = kids[0]["page"]
    # 페이지 역행(목차 파싱 오류)은 직전 값으로 보정
    last = 0
    for n in nodes:
        if n["page"] is None or n["page"] < last - 50:
            n["page"] = last
        last = max(last, n["page"])