RAG-CUT/backend/rag_cut/splitters/pdf_semantic.py

"""Generic heading/layout multimodal splitter for PDF-derived documents."""

from __future__ import annotations

import re
from dataclasses import dataclass, field

from rag_cut.models import Block, BlockType, Chunk, SplitConfig
from rag_cut.parsers.pdf.noise_filter import is_toc_noise_text, is_toc_title_text
from rag_cut.renderer import collect_chunk_layout_meta, render_blocks
from rag_cut.splitters.heading_splitter import (
    PURE_NUMBERED_HEADING_RE,
    TIME_LIKE_RE,
    looks_like_false_heading_text,
    looks_like_numbered_instruction,
    normalize_heading_block,
)

CHUNK_STRATEGY = "heading_layout_multimodal"

PAGE_NUMBER_RE = re.compile(r"^\s*(?:[-\u2013\u2014]?\s*)?\d{1,4}(?:\s*/\s*\d{1,4})?\s*$")
NUMBERED_HEADING_RE = re.compile(
    r"^\s*(?P<num>\d+(?:\.\d+)*)(?:[.．、:：]\s*|\s+)"
    r"(?P<title>[A-Za-z\u4e00-\u9fff][^\n]{1,120})\s*$"
)
LETTER_HEADING_RE = re.compile(
    r"^\s*(?P<letter>[A-Z])[\.)]\s+(?P<title>[A-Za-z0-9\u4e00-\u9fff][^\n]{1,100})\s*$"
)
# Procedural "Step N …" lines are body content, not section boundaries.
# Promoting them to headings caused heading-only groups to be dropped on flush
# (e.g. Steps 1–3 under 5.7.1 vanished while only Step 4 with following body survived).
STEP_INSTRUCTION_RE = re.compile(
    r"^\s*(?:Step\s*\d+|STEP\s*\d+|步骤\s*\d+)(?:\s*[:.：)\-]?\s+\S|\s*$)",
    re.I,
)
CHAPTER_HEADING_RE = re.compile(r"^\s*Chapter\s+\d+(?:\s*[:.-]?\s+[^\n]{1,100})?\s*$", re.I)
CN_HEADING_RE = re.compile(
    r"^\s*(?:[\u4e00-\u9fff]{1,3}[、.．]|[（(][\u4e00-\u9fff]{1,3}[)）])\s*[^\n]{1,100}$"
)


@dataclass
class _Group:
    heading_path: list[str] = field(default_factory=list)
    blocks: list[Block] = field(default_factory=list)


@dataclass
class _HeadingSignal:
    text: str
    level: int


def _text(block: Block) -> str:
    return (block.text or block.markdown or "").strip()


def _page(block: Block) -> int | None:
    page = block.meta.get("page")
    try:
        return int(page) if page is not None else None
    except (TypeError, ValueError):
        return None


def _bbox(block: Block) -> list[float]:
    bbox = block.meta.get("bbox")
    return list(bbox) if isinstance(bbox, (list, tuple)) else []


def _copy_with_meta(block: Block, **meta_updates) -> Block:
    meta = dict(block.meta)
    meta.update({k: v for k, v in meta_updates.items() if v not in (None, [], {})})
    return block.model_copy(update={"meta": meta})


def _is_empty_text_block(block: Block) -> bool:
    return block.type not in {BlockType.IMAGE, BlockType.TABLE} and not _text(block)


def _is_toc_text(text: str) -> bool:
    return is_toc_noise_text(text) or is_toc_title_text(text)


def _is_decorative_image(block: Block) -> bool:
    bbox = _bbox(block)
    if len(bbox) != 4:
        return False
    width = bbox[2] - bbox[0]
    height = bbox[3] - bbox[1]
    if width <= 0 or height <= 0:
        return True
    if width < 28 or height < 16:
        return True
    page_height = block.meta.get("page_height") or 0
    if page_height and bbox[3] <= page_height * 0.10 and width * height < 12000:
        return True
    return False


def _is_margin_noise(block: Block) -> bool:
    if block.type == BlockType.HEADING:
        return False
    text = _text(block)
    bbox = _bbox(block)
    page_height = block.meta.get("page_height") or 0
    if not bbox or not page_height:
        return False
    y0, y1 = bbox[1], bbox[3]
    in_top = y1 <= page_height * 0.08
    in_bottom = y0 >= page_height * 0.92
    if PAGE_NUMBER_RE.match(text) and (in_top or in_bottom):
        return True
    return len(text) <= 80 and (in_top or in_bottom) and block.meta.get("running_header")


def _running_header_texts(blocks: list[Block]) -> set[str]:
    page_count = len({_page(b) for b in blocks if _page(b) is not None})
    if page_count < 3:
        return set()
    locations: dict[str, set[int]] = {}
    for block in blocks:
        if block.type in {BlockType.IMAGE, BlockType.TABLE, BlockType.HEADING}:
            continue
        text = " ".join(_text(block).split())
        page = _page(block)
        bbox = _bbox(block)
        page_height = block.meta.get("page_height") or 0
        if not text or page is None or not bbox or not page_height or len(text) > 100:
            continue
        if bbox[3] <= page_height * 0.10 or bbox[1] >= page_height * 0.90:
            locations.setdefault(text, set()).add(page)
    threshold = max(3, int(page_count * 0.5))
    return {text for text, pages in locations.items() if len(pages) >= threshold}


def _is_noise(block: Block, running_headers: set[str]) -> bool:
    if _is_empty_text_block(block):
        return True
    text = " ".join(_text(block).split())
    if block.type == BlockType.IMAGE:
        return _is_decorative_image(block)
    if PAGE_NUMBER_RE.match(text) or TIME_LIKE_RE.match(text):
        return True
    if _is_toc_text(text):
        return True
    if text in running_headers:
        return True
    return _is_margin_noise(block)


def _heading_signal(block: Block, current_top_level: bool = False) -> _HeadingSignal | None:
    text = _text(block)
    if not text or len(text) > 180 or _is_toc_text(text):
        return None
    # Keep procedural steps inside the parent section; do not open a new group.
    if (
        STEP_INSTRUCTION_RE.match(text)
        or looks_like_numbered_instruction(text)
        or looks_like_false_heading_text(text)
    ):
        return None

    pure_numbered = PURE_NUMBERED_HEADING_RE.match(text)

    if block.type == BlockType.HEADING:
        level = block.level or 1
        numbered = NUMBERED_HEADING_RE.match(text)
        if pure_numbered:
            level = pure_numbered.group(1).count(".") + 1
        elif numbered:
            level = numbered.group("num").count(".") + 1
        elif LETTER_HEADING_RE.match(text):
            level = 2 if current_top_level else max(2, level)
        return _HeadingSignal(text=" ".join(text.split()), level=max(1, min(level, 6)))

    numbered = NUMBERED_HEADING_RE.match(text)
    if pure_numbered:
        return _HeadingSignal(text=" ".join(text.split()), level=pure_numbered.group(1).count(".") + 1)
    if numbered:
        return _HeadingSignal(text=" ".join(text.split()), level=numbered.group("num").count(".") + 1)
    if LETTER_HEADING_RE.match(text):
        return _HeadingSignal(text=" ".join(text.split()), level=2)
    if CHAPTER_HEADING_RE.match(text):
        return _HeadingSignal(text=" ".join(text.split()), level=1)
    if CN_HEADING_RE.match(text):
        return _HeadingSignal(text=" ".join(text.split()), level=2 if current_top_level else 1)

    font_size = block.meta.get("font_size")
    body_size = block.meta.get("body_font_size")
    if font_size and body_size and font_size >= body_size + 1.5 and len(text) <= 100:
        level = block.level or 2
        # Short CJK section banners (形態指標 / 策略指標) are chapter peers, not
        # subsections of the preceding numbered person/indicator entry.
        if (
            font_size >= body_size + 3
            and re.fullmatch(r"[\u4e00-\u9fffA-Za-z0-9/／\s]{2,24}", text)
            and not NUMBERED_HEADING_RE.match(text)
        ):
            level = 1
        return _HeadingSignal(text=" ".join(text.split()), level=max(1, min(level, 6)))
    return None


def _heading_block(block: Block, signal: _HeadingSignal) -> Block:
    meta = dict(block.meta)
    meta["section_boundary"] = True
    meta["heading_level"] = signal.level
    return Block(type=BlockType.HEADING, text=signal.text, level=signal.level, meta=meta)


def _enrich_block(block: Block, heading_path: list[str], group_blocks: list[Block]) -> Block:
    meta = dict(block.meta)
    meta["heading_path"] = list(heading_path)
    if heading_path:
        meta["nearest_heading"] = heading_path[-1]
        meta["parent_heading"] = heading_path[-1]
        meta["chapter"] = heading_path[0]
    meta.setdefault("section_boundary", block.type == BlockType.HEADING)
    if block.type == BlockType.HEADING:
        meta["heading_level"] = block.level or len(heading_path) or 1
    if block.type in {BlockType.IMAGE, BlockType.TABLE}:
        meta.setdefault("bound_heading", heading_path[-1] if heading_path else "")
        _bind_adjacent_text(meta, block, group_blocks)
    return block.model_copy(update={"meta": meta})


def _bind_adjacent_text(meta: dict, block: Block, group_blocks: list[Block]) -> None:
    page = _page(block)
    for prev in reversed(group_blocks):
        if prev.type == BlockType.PARAGRAPH and _page(prev) == page and _text(prev):
            meta.setdefault("preceding_text", _text(prev)[:400])
            break
        if prev.type == BlockType.HEADING:
            break
    if block.type == BlockType.TABLE and not meta.get("table_title"):
        for prev in reversed(group_blocks):
            if prev.type == BlockType.HEADING:
                meta["table_title"] = _text(prev)
                break
            if prev.type == BlockType.PARAGRAPH and _page(prev) == page and 0 < len(_text(prev)) <= 120:
                meta["table_title"] = _text(prev)
                break


def _fill_following_text(group: _Group) -> _Group:
    blocks = list(group.blocks)
    for i, block in enumerate(blocks):
        if block.type not in {BlockType.IMAGE, BlockType.TABLE}:
            continue
        meta = dict(block.meta)
        page = _page(block)
        for nxt in blocks[i + 1 :]:
            if nxt.type == BlockType.HEADING:
                break
            if nxt.type == BlockType.PARAGRAPH and _page(nxt) == page and _text(nxt):
                meta.setdefault("following_text", _text(nxt)[:400])
                break
        blocks[i] = block.model_copy(update={"meta": meta})
    return _Group(heading_path=group.heading_path, blocks=blocks)


def _group_has_body(group: _Group) -> bool:
    for block in group.blocks:
        if block.type in {BlockType.IMAGE, BlockType.TABLE}:
            return True
        if block.type != BlockType.HEADING and _text(block):
            return True
    # Safety net: long instructional "headings" are themselves searchable content.
    return any(block.type == BlockType.HEADING and len(_text(block)) >= 48 for block in group.blocks)


def _attach_heading_only_to_previous(groups: list[_Group], orphan: _Group) -> None:
    """Keep catalog-style heading-only groups instead of deleting them.

    Example: after \"共包含以下11個形態指標說明：\", lines like \"1.頭肩頂形態\"
    briefly open a group with no body before the next number arrives. Flush used
    to drop them; fold those leaf titles into the previous section as paragraphs.
    """
    if not groups or not orphan.blocks:
        return
    prev = groups[-1]
    prev_keys = {(b.type, _text(b), tuple(b.meta.get("bbox") or [])) for b in prev.blocks}
    path_set = set(prev.heading_path)
    for block in orphan.blocks:
        text = _text(block)
        if not text:
            continue
        # Skip copies of ancestor headings already injected into the orphan group.
        if text in path_set:
            continue
        key = (block.type, text, tuple(block.meta.get("bbox") or []))
        if key in prev_keys:
            continue
        if block.type == BlockType.HEADING:
            meta = dict(block.meta)
            meta["catalog_item"] = True
            meta["section_boundary"] = False
            demoted = Block(type=BlockType.PARAGRAPH, text=text, meta=meta)
            prev.blocks.append(_enrich_block(demoted, prev.heading_path, prev.blocks))
        else:
            prev.blocks.append(_enrich_block(block, prev.heading_path, prev.blocks))
        prev_keys.add(key)


def _group_is_retrievable(group: _Group) -> bool:
    """Return whether a rendered group contains searchable semantic context."""
    if group.heading_path:
        return True
    for block in group.blocks:
        if block.type == BlockType.TABLE:
            return True
        if block.type == BlockType.IMAGE:
            if block.ocr_text.strip():
                return True
            continue
        if block.type != BlockType.HEADING and _text(block):
            return True
    return False


def _render_group(group: _Group, index: int, config: SplitConfig, meta: dict | None = None) -> Chunk:
    group = _fill_following_text(group)
    layout_meta = collect_chunk_layout_meta(group.blocks)
    heading_path = group.heading_path or layout_meta.get("heading_path") or []
    heading = heading_path[-1] if heading_path else layout_meta.get("heading") or "Untitled section"
    extra_meta = {
        **layout_meta,
        "chunk_strategy": CHUNK_STRATEGY,
        "heading": heading,
        "heading_path": heading_path,
        "retrieval": _group_is_retrievable(group),
        **(meta or {}),
    }
    if "order_range" not in extra_meta:
        order_indices = [b.meta.get("order_index") for b in group.blocks if b.meta.get("order_index") is not None]
        if order_indices:
            extra_meta["order_range"] = [min(order_indices), max(order_indices)]
    chunk = render_blocks(group.blocks, index=index, meta=extra_meta)
    chunk_meta = dict(chunk.meta)
    chunk_meta.update(extra_meta)
    return chunk.model_copy(update={"meta": chunk_meta})


def _split_oversized_group(group: _Group, start_index: int, config: SplitConfig) -> list[Chunk]:
    """Split an oversized section by length only (no parent+child pair)."""
    rendered_len = sum(len(b.render()) + 2 for b in group.blocks)
    if rendered_len <= config.max_chunk_size or len(group.blocks) <= 3:
        return [_render_group(group, start_index, config)]

    heading_blocks = [b for b in group.blocks if b.type == BlockType.HEADING]
    prefix = heading_blocks[:1]
    chunks: list[Chunk] = []
    current: list[Block] = list(prefix)
    current_len = sum(len(b.render()) + 2 for b in current)
    part_index = 0

    def flush() -> None:
        nonlocal current, current_len, part_index
        body = [b for b in current if b.type != BlockType.HEADING]
        if not body:
            return
        part_group = _Group(group.heading_path, list(current))
        chunks.append(
            _render_group(
                part_group,
                start_index + len(chunks),
                config,
                meta={
                    "chunk_strategy": f"{CHUNK_STRATEGY}_part",
                    "part_index": part_index,
                    "retrieval": True,
                },
            )
        )
        part_index += 1
        current = list(prefix)
        current_len = sum(len(b.render()) + 2 for b in current)

    for block in group.blocks[len(prefix) :]:
        block_len = len(block.render()) + 2
        if current_len + block_len > config.max_chunk_size and len(current) > len(prefix):
            context_len = sum(
                len(item.render()) + 2
                for item in current
                if item.type != BlockType.HEADING
            )
            keep_atomic_context = (
                block.type in {BlockType.IMAGE, BlockType.TABLE}
                and context_len <= min(400, max(80, config.max_chunk_size // 2))
            )
            if not keep_atomic_context:
                flush()
        current.append(block)
        current_len += block_len
    flush()
    return chunks or [_render_group(group, start_index, config)]


def _build_groups(blocks: list[Block]) -> list[_Group]:
    running_headers = _running_header_texts(blocks)
    heading_stack: list[tuple[int, str, Block]] = []
    groups: list[_Group] = []
    current = _Group()

    def flush() -> None:
        nonlocal current
        if _group_has_body(current):
            groups.append(current)
        elif current.blocks:
            # Explicit parser headings with no body are empty document sections,
            # not catalog items to fold into the preceding section.
            if not any(block.meta.get("source_heading") for block in current.blocks):
                _attach_heading_only_to_previous(groups, current)
        current = _Group()

    for source in blocks:
        if _is_noise(source, running_headers):
            continue
        raw = normalize_heading_block(source)

        signal = _heading_signal(raw, current_top_level=bool(heading_stack))
        if signal:
            flush()
            while heading_stack and heading_stack[-1][0] >= signal.level:
                heading_stack.pop()
            heading = _heading_block(raw, signal)
            heading_stack.append((signal.level, signal.text, heading))
            path = [item[1] for item in heading_stack]
            current = _Group(heading_path=path)
            current.blocks.append(_enrich_block(heading, path, current.blocks))
            continue

        path = [item[1] for item in heading_stack]
        if not current.heading_path:
            current.heading_path = path

        current.blocks.append(_enrich_block(raw, current.heading_path, current.blocks))

    flush()
    return groups


def split_pdf_semantic(blocks: list[Block], config: SplitConfig) -> list[Chunk]:
    """Split PDF/Word-derived blocks into heading/layout-preserving multimodal chunks."""
    groups = _build_groups(blocks)
    chunks: list[Chunk] = []
    for group in groups:
        for chunk in _split_oversized_group(group, len(chunks), config):
            chunks.append(chunk.model_copy(update={"index": len(chunks)}))
    return chunks