app/ingest/parser_pdf.py

"""Layout- and table-aware PDF parsing (STEP 1).

Replaces the plain ``PyMuPDFLoader`` (which flattens tables into ambiguous
text and leaves OCR line-wrap corruption in place) with a hardened parser built
on PyMuPDF (``fitz``) — already a project dependency, so no new packages.

Per page we:

1. Extract text in reading order (``sort=True``) to reduce multi-column
   scrambling.
2. Detect tables with ``page.find_tables()`` and render each as Markdown
   wrapped in ``[TABLE]…[/TABLE]`` sentinels so condition-rating blocks and
   financial totals survive chunking with row/column structure intact and are
   NEVER flattened into prose.
3. Normalize the prose across pages: strip running headers/footers and page
   numbers, repair hyphenation, and unwrap hard-wrapped sentences.

Output remains a list of LangChain ``Document`` objects (one per page) with the
same ``page`` / ``source`` / ``total_pages`` metadata the rest of the pipeline
expects, so this is a drop-in replacement. If anything fails, we fall back to
the original ``PyMuPDFLoader`` so ingestion never hard-fails on a quirky PDF.
"""

from __future__ import annotations

import logging
from pathlib import Path

from langchain_core.documents import Document

from app.ingest.ocr_normalize import TABLE_CLOSE, TABLE_OPEN, normalize_pages

logger = logging.getLogger(__name__)


def _cell(value: object) -> str:
    """Render one table cell: strip, collapse newlines, escape pipes."""
    s = "" if value is None else str(value)
    return s.replace("\n", " ").replace("|", "\\|").strip()


def _table_to_markdown(rows: list[list[object]]) -> str:
    """Render extracted table rows as a GitHub-flavoured Markdown table."""
    cleaned = [[_cell(c) for c in row] for row in rows if row]
    cleaned = [r for r in cleaned if any(c for c in r)]
    if not cleaned:
        return ""
    width = max(len(r) for r in cleaned)
    cleaned = [r + [""] * (width - len(r)) for r in cleaned]
    header = cleaned[0]
    body = cleaned[1:]
    lines = ["| " + " | ".join(header) + " |", "| " + " | ".join(["---"] * width) + " |"]
    for r in body:
        lines.append("| " + " | ".join(r) + " |")
    return "\n".join(lines)


def _extract_page_tables(page: object) -> list[str]:
    """Return Markdown for each table on ``page`` (best-effort, never raises)."""
    out: list[str] = []
    try:
        finder = page.find_tables()
    except Exception as exc:  # find_tables can choke on malformed content
        logger.debug("find_tables failed on a page: %s", exc)
        return out
    tables = getattr(finder, "tables", None) or []
    for tbl in tables:
        try:
            md = _table_to_markdown(tbl.extract())
        except Exception as exc:  # noqa: BLE001 — a bad table must not kill the page
            logger.debug("table.extract failed: %s", exc)
            continue
        if md:
            out.append(f"{TABLE_OPEN}\n{md}\n{TABLE_CLOSE}")
    return out


def _parse_pdf_fitz(file_path: Path) -> list[Document]:
    """Primary path: layout/table-aware extraction via PyMuPDF."""
    import fitz  # PyMuPDF — already a dependency

    doc = fitz.open(str(file_path))
    try:
        total_pages = doc.page_count
        prose_pages: list[str] = []
        page_tables: list[list[str]] = []
        for page in doc:
            prose_pages.append(page.get_text("text", sort=True) or "")
            page_tables.append(_extract_page_tables(page))
    finally:
        doc.close()

    normalized = normalize_pages(prose_pages)

    documents: list[Document] = []
    for i, prose in enumerate(normalized):
        parts = [prose.strip()] if prose.strip() else []
        parts.extend(page_tables[i])
        content = "\n\n".join(parts)
        documents.append(
            Document(
                page_content=content,
                metadata={
                    "page": i,
                    "total_pages": total_pages,
                    "source": file_path.name,
                    "file_path": str(file_path),
                    "has_tables": bool(page_tables[i]),
                },
            )
        )
    return documents


def _parse_pdf_fallback(file_path: Path) -> list[Document]:
    """Fallback: original LangChain loader (no table structure, no normalization)."""
    from langchain_community.document_loaders import PyMuPDFLoader

    return PyMuPDFLoader(str(file_path)).load()


def parse_pdf(file_path: Path) -> list[Document]:
    """Load a ``.pdf`` and return one normalized, table-aware Document per page.

    Args:
        file_path: Absolute path to a ``.pdf`` file.

    Returns:
        List of :class:`langchain_core.documents.Document`, one per page, with
        prose normalized and tables preserved as Markdown inside
        ``[TABLE]…[/TABLE]`` sentinels.

    Raises:
        FileNotFoundError: If ``file_path`` does not exist.
        ValueError: If the file cannot be opened as a valid PDF by either path.
    """
    if not file_path.exists():
        raise FileNotFoundError(f"File not found: {file_path}")

    try:
        docs = _parse_pdf_fitz(file_path)
        logger.debug(
            "Parsed %d page(s) from %s (table-aware PyMuPDF)", len(docs), file_path.name
        )
        return docs
    except Exception as exc:  # noqa: BLE001 — degrade gracefully, never lose ingestion
        logger.warning(
            "Table-aware PDF parse failed for %s (%s); falling back to PyMuPDFLoader",
            file_path.name,
            exc,
        )
        try:
            return _parse_pdf_fallback(file_path)
        except Exception as exc2:
            raise ValueError(f"Cannot open PDF '{file_path}': {exc2}") from exc2