RAG-CUT/backend/rag_cut/parsers/pdf/layout.py

"""Page rendering and layout analysis."""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Any

import fitz

from rag_cut.parsers.pdf.tables import looks_like_table


_NUMBERED_LINE_RE = re.compile(
    r"^\s*(\d+(?:\.\d+)*)(?:\.|．)?\s*([A-Za-z0-9\u4e00-\u9fff].{2,})$"
)


@dataclass
class TextLine:
    x0: float
    y0: float
    x1: float
    y1: float
    text: str
    font_size: float


@dataclass
class LayoutRegion:
    x0: float
    y0: float
    x1: float
    y1: float
    kind: str  # text | table | image
    data: dict[str, Any] = field(default_factory=dict)


@dataclass
class PageLayout:
    page_index: int
    page_width: float
    page_height: float
    body_font_size: float
    regions: list[LayoutRegion] = field(default_factory=list)


def render_page(page: fitz.Page, zoom: float = 2.0) -> fitz.Pixmap:
    """Render page to pixmap for region cropping and OCR."""
    return page.get_pixmap(matrix=fitz.Matrix(zoom, zoom), alpha=False)


def _median_body_size(text_dict: dict) -> float:
    sizes: list[float] = []
    for block in text_dict.get("blocks", []):
        if block.get("type") != 0:
            continue
        for line in block.get("lines", []):
            for span in line.get("spans", []):
                sizes.append(span.get("size", 12))
    return sorted(sizes)[len(sizes) // 2] if sizes else 12


def _collect_text_lines(page: fitz.Page, body_size: float) -> list[TextLine]:
    text_dict = page.get_text("dict")
    lines: list[TextLine] = []
    for block in text_dict.get("blocks", []):
        if block.get("type") != 0:
            continue
        for line in block.get("lines", []):
            text = "".join(span.get("text", "") for span in line.get("spans", [])).strip()
            if not text:
                continue
            x0, y0, x1, y1 = line["bbox"]
            max_size = body_size
            for span in line.get("spans", []):
                max_size = max(max_size, span.get("size", body_size))
            lines.append(TextLine(x0=x0, y0=y0, x1=x1, y1=y1, text=text, font_size=max_size))
    return lines


def _detect_columns(
    items: list[tuple[float, float, float, float]],
    page_width: float,
) -> list[tuple[float, float, float, float]]:
    if not items:
        return items
    mid = page_width / 2
    left = [b for b in items if (b[0] + b[2]) / 2 < mid]
    right = [b for b in items if (b[0] + b[2]) / 2 >= mid]
    if len(left) >= 2 and len(right) >= 2:
        left.sort(key=lambda b: (b[1], b[0]))
        right.sort(key=lambda b: (b[1], b[0]))
        return left + right
    return sorted(items, key=lambda b: (b[1], b[0]))


def _sort_reading_order(
    items: list[tuple[float, float, float, float]],
    page_width: float,
) -> list[tuple[float, float, float, float]]:
    """Sort page boxes in a human reading order, preserving two-column flows."""
    if not items:
        return []

    mid = page_width / 2
    full_width: list[tuple[float, float, float, float]] = []
    column_items: list[tuple[float, float, float, float]] = []
    for box in items:
        width = box[2] - box[0]
        spans_mid = box[0] < mid < box[2]
        if width >= page_width * 0.60 or (spans_mid and width >= page_width * 0.35):
            full_width.append(box)
        else:
            column_items.append(box)

    left = [b for b in column_items if (b[0] + b[2]) / 2 < mid]
    right = [b for b in column_items if (b[0] + b[2]) / 2 >= mid]
    if len(left) < 2 or len(right) < 2:
        return sorted(items, key=lambda b: (b[1], b[0]))

    full_width.sort(key=lambda b: (b[1], b[0]))
    ordered: list[tuple[float, float, float, float]] = []
    segment_top = float("-inf")

    def add_columns_between(top: float, bottom: float) -> None:
        segment = [b for b in column_items if b[1] >= top and b[1] < bottom]
        segment_left = sorted([b for b in segment if (b[0] + b[2]) / 2 < mid], key=lambda b: (b[1], b[0]))
        segment_right = sorted([b for b in segment if (b[0] + b[2]) / 2 >= mid], key=lambda b: (b[1], b[0]))
        ordered.extend(segment_left)
        ordered.extend(segment_right)

    for box in full_width:
        add_columns_between(segment_top, box[1])
        ordered.append(box)
        segment_top = box[3]
    add_columns_between(segment_top, float("inf"))

    seen: set[tuple[float, float, float, float]] = set(ordered)
    ordered.extend(b for b in sorted(items, key=lambda b: (b[1], b[0])) if b not in seen)
    return ordered


def _rect_overlap(a: tuple[float, float, float, float], b: tuple[float, float, float, float]) -> float:
    x0 = max(a[0], b[0])
    y0 = max(a[1], b[1])
    x1 = min(a[2], b[2])
    y1 = min(a[3], b[3])
    if x1 <= x0 or y1 <= y0:
        return 0.0
    inter = (x1 - x0) * (y1 - y0)
    area_a = max((a[2] - a[0]) * (a[3] - a[1]), 1e-6)
    return inter / area_a


def _center_inside(inner: tuple[float, float, float, float], outer: tuple[float, float, float, float]) -> bool:
    cx = (inner[0] + inner[2]) / 2
    cy = (inner[1] + inner[3]) / 2
    return outer[0] <= cx <= outer[2] and outer[1] <= cy <= outer[3]


def _horizontal_overlap_ratio(a: TextLine, b: TextLine) -> float:
    overlap = min(a.x1, b.x1) - max(a.x0, b.x0)
    if overlap <= 0:
        return 0.0
    return overlap / max(min(a.x1 - a.x0, b.x1 - b.x0), 1e-6)


def _line_on_image(
    line: TextLine,
    image_boxes: list[tuple[float, float, float, float]],
    large_image_boxes: list[tuple[float, float, float, float]],
) -> bool:
    # Never drop numbered section titles even if they sit near a screenshot.
    if _NUMBERED_LINE_RE.match(line.text.strip()):
        return False
    bbox = (line.x0, line.y0, line.x1, line.y1)
    # Large background/screenshot: require near-full coverage of the line, not
    # merely center-inside (which wiped text sitting in margins of wide figures).
    if large_image_boxes and any(
        _center_inside(bbox, box) and _rect_overlap(bbox, box) >= 0.85 for box in large_image_boxes
    ):
        return True
    return any(_rect_overlap(bbox, box) >= 0.55 for box in image_boxes)


def _area(box: tuple[float, float, float, float]) -> float:
    return max(box[2] - box[0], 0.0) * max(box[3] - box[1], 0.0)


def _filter_nested_image_regions(regions: list[LayoutRegion]) -> list[LayoutRegion]:
    """Drop image fragments that are already contained in a larger screenshot."""
    result: list[LayoutRegion] = []
    boxes = [(r.x0, r.y0, r.x1, r.y1) for r in regions]
    for region, box in zip(regions, boxes):
        box_area = _area(box)
        nested = False
        for other in boxes:
            other_area = _area(other)
            if other == box or other_area <= box_area * 1.5:
                continue
            if _center_inside(box, other) and _rect_overlap(box, other) >= 0.85:
                nested = True
                break
        if not nested:
            result.append(region)
    return result


def _merge_lines_to_paragraphs(lines: list[TextLine], page_width: float) -> list[TextLine]:
    if not lines:
        return []

    ordered_boxes = _sort_reading_order([(ln.x0, ln.y0, ln.x1, ln.y1) for ln in lines], page_width)
    order = {(b[0], b[1], b[2], b[3]): i for i, b in enumerate(ordered_boxes)}
    ordered = sorted(lines, key=lambda ln: order.get((ln.x0, ln.y0, ln.x1, ln.y1), (ln.y0, ln.x0)))

    paragraphs: list[TextLine] = []
    current = ordered[0]
    for nxt in ordered[1:]:
        vgap = nxt.y0 - current.y1
        line_h = max(current.y1 - current.y0, nxt.y1 - nxt.y0, 8.0)
        size_gap = abs(current.font_size - nxt.font_size)
        size_ratio = max(current.font_size, nxt.font_size) / max(
            min(current.font_size, nxt.font_size), 1e-6
        )
        # Keep title vs body separate: small absolute gap is enough when ratio is large.
        same_style = size_gap <= 1.0 and size_ratio <= 1.15
        either_numbered = bool(
            _NUMBERED_LINE_RE.match(current.text.strip()) or _NUMBERED_LINE_RE.match(nxt.text.strip())
        )
        if (
            same_style
            and not either_numbered
            and vgap <= line_h * 2.2
            and _horizontal_overlap_ratio(current, nxt) >= 0.25
        ):
            joiner = "" if current.text.endswith("-") or current.text.endswith(" ") else " "
            current = TextLine(
                x0=min(current.x0, nxt.x0),
                y0=current.y0,
                x1=max(current.x1, nxt.x1),
                y1=nxt.y1,
                text=f"{current.text}{joiner}{nxt.text}",
                font_size=max(current.font_size, nxt.font_size),
            )
        else:
            paragraphs.append(current)
            current = nxt
    paragraphs.append(current)
    return paragraphs


def analyze_page_layout(page: fitz.Page, page_index: int) -> PageLayout:
    """Layout analysis: text paragraphs, table regions, image regions."""
    text_dict = page.get_text("dict")
    body_size = _median_body_size(text_dict)
    page_width = page.rect.width
    page_height = page.rect.height
    page_area = page_width * page_height

    table_regions: list[LayoutRegion] = []
    try:
        table_finder = page.find_tables()
        for idx, table in enumerate(table_finder.tables):
            bbox = table.bbox
            rows = table.extract() or []
            if not looks_like_table(rows):
                continue
            table_regions.append(
                LayoutRegion(
                    x0=bbox[0],
                    y0=bbox[1],
                    x1=bbox[2],
                    y1=bbox[3],
                    kind="table",
                    data={"rows": rows, "table_index": idx, "source": "pymupdf"},
                )
            )
    except Exception:
        pass

    image_regions: list[LayoutRegion] = []
    for img_info in page.get_images(full=True):
        xref = img_info[0]
        try:
            rects = page.get_image_rects(xref)
        except Exception:
            rects = []
        if not rects:
            continue
        for rect_idx, rect in enumerate(rects):
            image_regions.append(
                LayoutRegion(
                    x0=rect.x0,
                    y0=rect.y0,
                    x1=rect.x1,
                    y1=rect.y1,
                    kind="image",
                    data={"xref": xref, "rect_index": rect_idx},
                )
            )
    image_regions = _filter_nested_image_regions(image_regions)

    image_boxes = [(r.x0, r.y0, r.x1, r.y1) for r in image_regions]
    large_image_boxes = [
        box
        for box in image_boxes
        if (box[2] - box[0]) * (box[3] - box[1]) >= page_area * 0.12
    ]
    table_boxes = [(r.x0, r.y0, r.x1, r.y1) for r in table_regions]

    raw_lines = _collect_text_lines(page, body_size)
    filtered_lines = [
        ln
        for ln in raw_lines
        if not _line_on_image(ln, image_boxes, large_image_boxes)
        and not any(_rect_overlap((ln.x0, ln.y0, ln.x1, ln.y1), box) >= 0.55 for box in table_boxes)
    ]
    paragraphs = _merge_lines_to_paragraphs(filtered_lines, page_width)

    text_regions: list[LayoutRegion] = []
    for para in paragraphs:
        text_regions.append(
            LayoutRegion(
                x0=para.x0,
                y0=para.y0,
                x1=para.x1,
                y1=para.y1,
                kind="text",
                data={"text": para.text, "font_size": para.font_size},
            )
        )

    regions = text_regions + table_regions + image_regions
    ordered_boxes = _sort_reading_order(
        [(r.x0, r.y0, r.x1, r.y1) for r in regions],
        page_width,
    )
    order = {box: i for i, box in enumerate(ordered_boxes)}
    regions.sort(key=lambda r: order.get((r.x0, r.y0, r.x1, r.y1), len(order)))

    return PageLayout(
        page_index=page_index,
        page_width=page_width,
        page_height=page_height,
        body_font_size=body_size,
        regions=regions,
    )