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

"""Cross-page table detection and merging."""

from __future__ import annotations

from rag_cut.models import Block, BlockType
from rag_cut.parsers.pdf.tables import (
    detect_header_row_count,
    header_signature,
    normalize_rows,
    rows_to_markdown,
    split_body_and_footnotes,
)


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


def _header_rows(block: Block) -> int:
    rows = block.meta.get("rows") or []
    return int(block.meta.get("header_rows") or detect_header_row_count(normalize_rows(rows)) or 1)


def _header_sig(block: Block) -> tuple[tuple[str, ...], ...]:
    if block.meta.get("header_signature"):
        raw = block.meta["header_signature"]
        return tuple(tuple(r) for r in raw)
    rows = normalize_rows(block.meta.get("rows") or [])
    return header_signature(rows, _header_rows(block))


def _column_count(block: Block) -> int:
    rows = normalize_rows(block.meta.get("rows") or [])
    return max((len(r) for r in rows), default=0)


def _similar_columns(a: Block, b: Block) -> bool:
    ca, cb = _column_count(a), _column_count(b)
    if ca == 0 or cb == 0:
        return False
    return ca == cb or abs(ca - cb) <= 1


def _repeated_header(rows: list[list[str]], sig: tuple[tuple[str, ...], ...]) -> int:
    """Return number of leading rows in `rows` that repeat the header signature."""
    if not sig:
        return 0
    n = len(sig)
    if len(rows) < n:
        return 0
    if header_signature(rows, n) == sig:
        return n
    if n == 1 and rows and tuple(rows[0]) == sig[0]:
        return 1
    return 0


def _can_merge_continuation(prev: Block, nxt: Block) -> bool:
    if prev.type != BlockType.TABLE or nxt.type != BlockType.TABLE:
        return False

    prev_page = _page(prev)
    nxt_page = _page(nxt)
    if prev_page is None or nxt_page is None or nxt_page != prev_page + 1:
        return False
    if not _similar_columns(prev, nxt):
        return False

    sig = _header_sig(prev)
    if not sig:
        return False

    rows_b = normalize_rows(nxt.meta.get("rows") or [])
    if not rows_b:
        return False

    if _repeated_header(rows_b, sig) > 0:
        return True

    # Continuation without repeated header: similar width and no title on next table
    if nxt.meta.get("table_title"):
        return False

    prev_bbox = prev.meta.get("bbox") or []
    nxt_bbox = nxt.meta.get("bbox") or []
    if len(prev_bbox) == 4 and len(nxt_bbox) == 4:
        prev_width = prev_bbox[2] - prev_bbox[0]
        nxt_width = nxt_bbox[2] - nxt_bbox[0]
        if prev_width > 0 and abs(prev_width - nxt_width) / prev_width <= 0.15:
            return True

    return _column_count(prev) == _column_count(nxt)


def _merge_two_tables(prev: Block, nxt: Block) -> Block:
    rows_a = normalize_rows(prev.meta.get("rows") or [])
    rows_b = normalize_rows(nxt.meta.get("rows") or [])
    header_rows = _header_rows(prev)
    sig = _header_sig(prev)

    skip = _repeated_header(rows_b, sig)
    merged_rows = rows_a + rows_b[skip:]

    body_rows, _, foot_a = split_body_and_footnotes(rows_a, header_rows)
    _, _, foot_b = split_body_and_footnotes(rows_b, skip or header_rows)
    footnotes = " ".join(x for x in (prev.meta.get("footnotes") or foot_a, foot_b) if x).strip()

    pages = sorted(set((prev.meta.get("pages") or [_page(prev)]) + [_page(nxt)]))
    pages = [p for p in pages if p is not None]
    bboxes = list(prev.meta.get("bboxes") or [])
    if prev.meta.get("bbox"):
        bboxes.append({"page": _page(prev), "bbox": prev.meta["bbox"]})
    if nxt.meta.get("bbox"):
        bboxes.append({"page": _page(nxt), "bbox": nxt.meta["bbox"]})

    prev_bbox = prev.meta.get("bbox") or [0, 0, 0, 0]
    nxt_bbox = nxt.meta.get("bbox") or prev_bbox
    merged_bbox = [
        min(prev_bbox[0], nxt_bbox[0]),
        min(prev_bbox[1], nxt_bbox[1]),
        max(prev_bbox[2], nxt_bbox[2]),
        max(prev_bbox[3], nxt_bbox[3]),
    ]

    md = rows_to_markdown(merged_rows, header_rows=header_rows, include_footnotes=footnotes)
    meta = dict(prev.meta)
    meta.update(
        {
            "rows": merged_rows,
            "row_count": len(merged_rows),
            "col_count": max((len(r) for r in merged_rows), default=0),
            "header_rows": header_rows,
            "header_signature": [list(r) for r in sig],
            "footnotes": footnotes,
            "pages": pages,
            "page": pages[0] if pages else prev.meta.get("page"),
            "bbox": merged_bbox,
            "bboxes": bboxes,
            "cross_page": len(pages) > 1,
            "merged_table_count": int(prev.meta.get("merged_table_count") or 1) + 1,
        }
    )
    if nxt.meta.get("following_text") and not meta.get("following_text"):
        meta["following_text"] = nxt.meta.get("following_text")

    return prev.model_copy(update={"markdown": md, "meta": meta})


def merge_cross_page_tables(blocks: list[Block]) -> list[Block]:
    """Merge consecutive cross-page table blocks that share headers/structure."""
    if not blocks:
        return []

    result: list[Block] = []
    i = 0
    while i < len(blocks):
        current = blocks[i]
        if current.type != BlockType.TABLE:
            result.append(current)
            i += 1
            continue

        merged = current
        j = i + 1
        while j < len(blocks):
            nxt = blocks[j]
            if nxt.type == BlockType.TABLE and _can_merge_continuation(merged, nxt):
                merged = _merge_two_tables(merged, nxt)
                j += 1
                continue
            break

        result.append(merged)
        i = j

    return result