brain-of-reese/app/rag/pdf.py

"""PDF parsing utilities (phase 126 — PDF support, locked decision A1).

Pure functions (I/O = reading the PDF file only) for extracting text
and embedded images from PDF files using PyMuPDF (the ``pymupdf``
package, imported as ``fitz`` — its stable API name).

* :func:`extract_pdf_text` — the text of every page that has text,
  each prefixed with :data:`PAGE_SEPARATOR` (``---PAGE-N---``, N is the
  1-indexed page number). The PDF chunker (phase 126 task 03) splits on
  these markers so every chunk keeps its page anchor. Pages without
  text (image-only pages) contribute nothing, so a PDF with no
  extractable text returns ``""``.
* :func:`extract_pdf_page_texts` — one entry per page that has text:
  1-indexed page number → the page's stripped text (the same page read
  as :func:`extract_pdf_text`, so the two agree page-for-page; the
  value for page N is exactly that page's separator-block body). The
  phase-137 image-context block is built from this map (page N's entry
  is the document context for an image embedded on page N).
* :func:`extract_pdf_images` — one :class:`PdfImage` per embedded-image
  occurrence, in page order then 0-based index on the page. ``data`` is
  the image's original encoded bytes as stored in the PDF (no
  re-encoding), ``mime`` is derived via :func:`_mime_from_image_info`.
* :func:`_mime_from_image_info` — maps a PyMuPDF image-info dict to a
  MIME type: the embedded file extension (``ext``) first, then the
  image type code (``type``: 2=JPEG, 3=PNG, 4=GIF, 8=TIFF, 9=WebP),
  else the ``application/octet-stream`` fallback.
* :func:`is_usable_pdf_text` — the phase-133 garbage gate: a pure,
  deterministic predicate over the extracted text layer (no PDF, no
  I/O) — REAL content vs. garbage (the pre-126 XML/binary soup, the
  empty text layer). It is the single trigger for the page-screenshot
  vision fallback (phase 133 task 03).
* :func:`pictorial_pages` — the phase-140 sparse-page gate: the 1-based
  page numbers of a :func:`extract_pdf_page_texts` map whose stripped
  text is STRICTLY under *min_chars* chars (ascending) — pure,
  deterministic, NO LLM call (LOCKED A1): the "way to determine when
  that's necessary" (TODO L5) for the pictorial pass.
* :func:`render_pdf_pages` — the fallback's visual input (phase 133
  task 03): one PNG per page (150 DPI, page order); the phase-140
  optional 1-based *pages* selector renders ONLY the named pages
  (ascending — a 500-page PDF with 30 sparse pages renders 30 PNGs,
  not 500), ``None`` = every page (the phase-133 caller's result is
  byte-identical); an unreadable (or out-of-range named) page yields
  ``b""`` at its position (the caller's placeholder — the result
  stays aligned with the requested pages), an unopenable file yields
  ``[]`` (the caller's existing empty-content path).

Every failure mode is fail-soft: an unreadable file returns an empty
result, an unreadable page or image is logged and skipped — this module
never raises for malformed input (the importer owns the loud failures).
"""
from __future__ import annotations

import logging
import re
from collections.abc import Iterator, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import pymupdf as fitz  # PyMuPDF (locked A1); `fitz` is the stable API name

logger = logging.getLogger("app.rag.pdf")

#: Page separator used in extracted text (the chunker splits on this).
#: ``{page}`` is the 1-indexed page number.
PAGE_SEPARATOR = "---PAGE-{page}---"

#: Mapping from PyMuPDF image type codes to MIME types
#: (2=JPEG, 3=PNG, 4=GIF, 8=TIFF, 9=WebP — phase-126 design).
_TYPE_TO_MIME: dict[int, str] = {
    2: "image/jpeg",
    3: "image/png",
    4: "image/gif",
    8: "image/tiff",
    9: "image/webp",
}

#: Mapping from embedded-image file extensions to MIME types.
_EXT_TO_MIME: dict[str, str] = {
    "png": "image/png",
    "jpg": "image/jpeg",
    "jpeg": "image/jpeg",
    "gif": "image/gif",
    "webp": "image/webp",
    "tif": "image/tiff",
    "tiff": "image/tiff",
}


@dataclass
class PdfImage:
    """One embedded image from a PDF page."""

    page: int  #: 1-indexed page number
    index: int  #: image index on the page (0-based, occurrence order)
    width: int  #: native pixel width
    height: int  #: native pixel height
    data: bytes  #: original encoded bytes as stored in the PDF
    mime: str  #: MIME type (e.g. ``image/png``)


def _mime_from_image_info(info: dict[str, Any]) -> str:
    """Map a PyMuPDF image-info dict to a MIME type.

    The embedded file extension (``ext``) wins over the image type code
    (``type``); unknown/missing both fall back to
    ``application/octet-stream``.
    """
    ext = str(info.get("ext") or "").lower()
    if ext in _EXT_TO_MIME:
        return _EXT_TO_MIME[ext]
    type_code = info.get("type")
    if isinstance(type_code, int):
        return _TYPE_TO_MIME.get(type_code, "application/octet-stream")
    return "application/octet-stream"


def _page_image_entries(page: Any) -> list[dict[str, Any]]:
    """One info dict per embedded-image occurrence on *page*.

    Prefers ``get_image_info(xrefs=True)`` (one entry per occurrence,
    with the ``xref`` object number); older PyMuPDF releases without the
    ``xrefs`` keyword are served from ``get_images(full=True)`` (one
    entry per unique image).
    """
    try:
        infos: list[dict[str, Any]] = page.get_image_info(xrefs=True)
    except TypeError:
        infos = [
            {"xref": entry[0], "width": entry[2], "height": entry[3]}
            for entry in page.get_images(full=True)
        ]
    return [dict(info) for info in infos]


def _iter_page_texts(doc: Any, pdf_path: Path) -> Iterator[tuple[int, str]]:
    """Yield ``(page, text)`` for every page of *doc* that has text.

    ``page`` is the 1-indexed page number and ``text`` the page's
    stripped text (the ``get_text("text")`` read — it always returns a
    str; the union return type comes from pymupdf's other flavors,
    dict/words/blocks). A page that cannot be read is logged and
    skipped; a page without text contributes nothing. The shared page
    iteration behind :func:`extract_pdf_text` and
    :func:`extract_pdf_page_texts` — the two callers share the exact
    same read and cannot drift.
    """
    for page_num in range(len(doc)):
        try:
            raw = doc[page_num].get_text("text")
        except Exception as exc:
            logger.warning(
                "pdf: could not read text of page %d of %s — %s",
                page_num + 1, pdf_path, exc,
            )
            continue
        text = raw.strip() if isinstance(raw, str) else ""
        if text:
            yield page_num + 1, text


def extract_pdf_text(pdf_path: Path) -> str:
    """Extract the text of all pages of *pdf_path*.

    Every page that has text is prefixed with :data:`PAGE_SEPARATOR`
    (``---PAGE-N---``, N 1-indexed) and the page blocks are joined with
    a blank line. Pages without text contribute nothing, so a PDF with
    no extractable text (or a zero-page file) returns ``""``.

    Fail-soft: an unreadable file returns ``""``; a page that cannot be
    read is logged and skipped.
    """
    try:
        doc = fitz.open(str(pdf_path))
    except Exception as exc:
        logger.warning("pdf: could not open %s — %s", pdf_path, exc)
        return ""

    with doc:
        parts = [
            f"{PAGE_SEPARATOR.format(page=page)}\n{text}"
            for page, text in _iter_page_texts(doc, pdf_path)
        ]
    return "\n\n".join(parts)


def extract_pdf_page_texts(pdf_path: Path) -> dict[int, str]:
    """One entry per page of *pdf_path* that has text.

    Keys are the 1-indexed page numbers, values the page's stripped
    text — the shared :func:`_iter_page_texts` read, so this map agrees
    with :func:`extract_pdf_text` page-for-page: the value for page N
    is exactly the text after that page's ``---PAGE-N---`` separator.
    Pages without text (image-only pages) contribute nothing — no key.

    The phase-137 document context for an embedded image is built from
    this map: one image on page N gets page N's text as the context
    block's body.

    Fail-soft: an unreadable file returns ``{}`` (one warning — the
    :func:`extract_pdf_text` open contract); a page that cannot be
    read is logged and skipped.
    """
    try:
        doc = fitz.open(str(pdf_path))
    except Exception as exc:
        logger.warning("pdf: could not open %s for page texts — %s", pdf_path, exc)
        return {}

    with doc:
        return dict(_iter_page_texts(doc, pdf_path))


#: A line "looks like XML" when it carries a tag-shaped token — the
#: pre-126 garbage's signature (raw XMP blocks, PDF-structure noise).
_XMLISH_LINE_RE = re.compile(r"</?[A-Za-z_][\w:.-]*[\s/>=]")

#: Garbage-gate thresholds (phase 133 task 02): an extracted text
#: layer is UNUSABLE when its XML-ish line fraction reaches this bar.
#: The FINDING's reproduced pre-126 read of the owner's file is 46%
#: XML-ish; a real document ABOUT XML that reaches the bar is a
#: documented false positive (ASSUMPTION A8 — the fallback still
#: extracts it via vision; the cost is a few vision calls).
XMLISH_LINE_FRACTION = 0.30

#: Garbage-gate thresholds (phase 133 task 02): the non-printable
#: fraction (U+FFFD replacement chars + C0 controls over the whole
#: text) at which the layer is UNUSABLE — the ``errors="replace"``
#: binary-soup signal (the pre-126 read is ~19% non-printable; real
#: extracted text is ~0.05%).
NONPRINT_CHAR_FRACTION = 0.05


def is_usable_pdf_text(text: str) -> bool:
    """Whether an extracted PDF text layer is REAL content (not garbage).

    Unusable (→ the page-screenshot fallback) when ANY of:

    * ``text.strip() == ""`` — the empty text layer (a scanned/image-
      only PDF has nothing to read);
    * ≥ :data:`XMLISH_LINE_FRACTION` (0.30) of the non-empty lines
      match :data:`_XMLISH_LINE_RE` — the raw-XMP / binary-soup
      signature (the FINDING's reproduced pre-126 read is 46% —
      ``read_text(errors="replace")`` on the owner's file left
      ``<illustrator:StartupProfile>Web</illustrator:StartupProfile>``
      lines in the preview);
    * ≥ :data:`NONPRINT_CHAR_FRACTION` (0.05) of all characters are
      non-printable — the ``errors="replace"`` binary-soup signal.
      The non-printable set is ``U+FFFD`` + C0 control chars
      (0x00–0x1F) EXCEPT the structural ``\\t``/``\\n``/``\\r`` the
      extractor legitimately emits (the explicit set is deliberate: a
      ``unicodedata`` Cc-category check would catch those too, and
      counting them would flag every prose document).

    Pure + deterministic (no PDF, no I/O) — the unit tests pin the
    thresholds against the REAL reproduced garbage string, not a
    hand-drawn one. A real document ABOUT XML whose tag-like lines
    reach 30% is a documented false positive (ASSUMPTION A8 — the
    fallback still extracts it correctly via vision; the cost is a
    few vision calls).
    """
    if not text.strip():
        return False
    non_empty_lines = [line for line in text.splitlines() if line.strip()]
    xml_ish = sum(1 for line in non_empty_lines if _XMLISH_LINE_RE.search(line))
    if xml_ish / len(non_empty_lines) >= XMLISH_LINE_FRACTION:
        return False
    non_printable = sum(
        1
        for ch in text
        if ch == "\ufffd" or (ord(ch) < 32 and ch not in "\t\n\r")
    )
    return non_printable / len(text) < NONPRINT_CHAR_FRACTION


def pictorial_pages(page_texts: dict[int, str], *, min_chars: int) -> list[int]:
    """The 1-based page numbers of *page_texts* that are PICTORIAL
    candidates: ``len(text.strip()) < min_chars`` — STRICTLY less than
    (a page with exactly *min_chars* stripped chars is NOT sparse; the
    boundary is pinned in the unit tests), ascending.

    The deterministic sparse-page gate (phase 140, task 01, LOCKED A1)
    — the "way to determine when that's necessary" for the pictorial PDF
    pass (TODO L5): NO LLM call, pure, deterministic (no PDF, no I/O —
    the input is a :func:`extract_pdf_page_texts` map, so the gate never
    re-opens the file). A page that carries no text at all has NO key in
    the map and is therefore not a candidate here (the importer's pass
    owns the page-total handling); an empty map yields ``[]`` — a
    clean-text PDF's path never burns a vision call.

    The IKEA example (the phase-140 fixture, ``min_chars=100``):
    ``{1: "BESTÅ", 2: <7 KB of warnings>, 6: "Illustration 6"}`` →
    ``[1, 6, …]`` — page 1 (the title page: the logo + product name) IS
    a candidate: describing it is useful (the logo + product name) and
    never harmful; page 2's 7 KB of multilingual warnings is not.

    ``min_chars=0`` is the documented PLAN A21 special case — **every
    page in the map is a candidate** (the "describe everything" mode).
    The strict-``<`` test cannot express this (no length is below 0), so
    0 is not "a stricter bar": it is the A21 no-cap endpoint, and a
    negative value — unreachable through ``Settings`` (the validator
    fails startup loudly) — takes the same branch defensively.
    """
    if min_chars <= 0:
        return sorted(page_texts)
    return sorted(
        page for page, text in page_texts.items() if len(text.strip()) < min_chars
    )


def render_pdf_pages(
    pdf_path: Path, dpi: int = 150, pages: Sequence[int] | None = None
) -> list[bytes]:
    """Render the selected pages of *pdf_path* to PNGs.

    *pages* (phase 140, task 02): the optional 1-based selector —
    render ONLY the named pages, in ascending order, one PNG each
    (a 500-page PDF with 30 pictorial candidates renders 30 PNGs, not
    500). ``None`` (the default) = every page in page order — the
    phase-133 caller's result is byte-identical, so list index *i* is
    page *i*+1 (with a selector, list index *i* is the *i*-th smallest
    NAMED page).

    The fallback's visual input (phase 133 task 03, ASSUMPTION A8:
    fixed 150 DPI — the balance of vision legibility vs data-URL
    payload). A page that cannot be rendered — including a named page
    outside the document's page range — yields ``b""`` at ITS position
    (the caller's placeholder — the result's order NEVER shifts, so it
    stays aligned with the requested pages; the failure is logged, the
    :func:`extract_pdf_text` fail-soft style). A file that cannot be
    OPENED yields ``[]`` (the caller's existing empty-content path —
    the doc lands with the no-extractable-text behavior, unchanged).
    """
    try:
        doc = fitz.open(str(pdf_path))
    except Exception as exc:
        logger.warning("pdf: could not open %s for page rendering — %s", pdf_path, exc)
        return []

    # 1-based page numbers: the default walks every page in page order
    # (the phase-133 contract — list index *i* is page *i*+1); a
    # selector renders the named pages, ascending (de-duplicated — a
    # repeated page number would render it twice for no reason).
    numbers = (
        range(1, len(doc) + 1)
        if pages is None
        else sorted(set(pages))
    )
    rendered: list[bytes] = []
    with doc:
        for page_number in numbers:
            try:
                rendered.append(
                    doc[page_number - 1].get_pixmap(dpi=dpi).tobytes("png")
                )
            except Exception as exc:
                logger.warning(
                    "pdf: could not render page %d of %s — %s",
                    page_number, pdf_path, exc,
                )
                rendered.append(b"")
    return rendered


def extract_pdf_images(pdf_path: Path) -> list[PdfImage]:
    """Extract all embedded images from *pdf_path*.

    Returns one :class:`PdfImage` per image occurrence, in page order
    then 0-based index on the page. ``data`` is the original encoded
    bytes as stored in the PDF; ``mime`` is derived from the image's
    extension/type via :func:`_mime_from_image_info`.

    Fail-soft: an unreadable file returns ``[]``; an image without an
    ``xref`` or that cannot be extracted is logged and skipped.
    """
    try:
        doc = fitz.open(str(pdf_path))
    except Exception as exc:
        logger.warning("pdf: could not open %s for image extraction — %s", pdf_path, exc)
        return []

    images: list[PdfImage] = []
    with doc:
        for page_num in range(len(doc)):
            try:
                entries = _page_image_entries(doc[page_num])
            except Exception as exc:
                logger.warning(
                    "pdf: could not list images on page %d of %s — %s",
                    page_num + 1, pdf_path, exc,
                )
                continue
            for index, info in enumerate(entries):
                xref = int(info.get("xref") or 0)
                if not xref:
                    logger.warning(
                        "pdf: skipping image page=%d index=%d of %s — no xref",
                        page_num + 1, index, pdf_path,
                    )
                    continue
                try:
                    extracted = doc.extract_image(xref)
                except Exception as exc:
                    logger.warning(
                        "pdf: could not extract image page=%d index=%d of %s — %s",
                        page_num + 1, index, pdf_path, exc,
                    )
                    continue
                data = extracted.get("image") or b""
                if not data:
                    logger.warning(
                        "pdf: skipping image page=%d index=%d of %s — empty data",
                        page_num + 1, index, pdf_path,
                    )
                    continue
                # The image info's own `ext` is unreliable in some
                # PyMuPDF versions (None); `extract_image` reports the
                # actual format of the returned bytes, so it wins.
                mime = _mime_from_image_info({**info, "ext": extracted.get("ext") or ""})
                images.append(
                    PdfImage(
                        page=page_num + 1,
                        index=index,
                        width=int(extracted.get("width") or info.get("width") or 0),
                        height=int(extracted.get("height") or info.get("height") or 0),
                        data=data,
                        mime=mime,
                    )
                )
    return images