backend/app/ocr/layout.py

"""Stage 1a — digital text layer + page rasterization.

Uses PyMuPDF (fitz) for the embedded text layer and word bounding boxes, and
pdfplumber (if installed) for higher-fidelity table-aware extraction. This is the
fast, exact, free channel — no model required. For native digital PDFs it is all
you need; the OCR channel only earns its keep on scans/photos.
"""
from __future__ import annotations

import importlib.util
from dataclasses import dataclass, field
from pathlib import Path

IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"}


@dataclass
class Block:
    text: str
    page: int
    bbox: tuple[float, float, float, float] | None = None
    source: str = "text"  # text | ocr | fused
    confidence: float = 1.0


@dataclass
class ChannelResult:
    text: str = ""
    blocks: list[Block] = field(default_factory=list)
    pages: int = 0
    available: bool = False
    engine: str = "none"

    @property
    def char_count(self) -> int:
        return len(self.text.strip())


def _has(mod: str) -> bool:
    return importlib.util.find_spec(mod) is not None


def extract_text_layer(path: str | Path) -> ChannelResult:
    """Extract the embedded text layer from a PDF (empty for scanned PDFs/images)."""
    path = Path(path)
    if path.suffix.lower() in IMAGE_EXTS:
        return ChannelResult(available=False, engine="none", pages=1)
    if not _has("fitz"):
        return ChannelResult(available=False, engine="none")

    import fitz  # PyMuPDF

    blocks: list[Block] = []
    parts: list[str] = []
    try:
        doc = fitz.open(str(path))
    except Exception:
        return ChannelResult(available=False, engine="none")
    for pno in range(doc.page_count):
        page = doc.load_page(pno)
        page_text = page.get_text("text")
        parts.append(page_text)
        # word-level boxes for the provenance overlay
        for w in page.get_text("words"):
            x0, y0, x1, y1, word = w[0], w[1], w[2], w[3], w[4]
            blocks.append(Block(text=word, page=pno, bbox=(x0, y0, x1, y1), source="text"))
    doc.close()
    text = "\n".join(parts)
    return ChannelResult(
        text=text, blocks=blocks, pages=len(parts),
        available=len(text.strip()) > 0, engine="pymupdf",
    )


def page_count(path: str | Path) -> int:
    path = Path(path)
    if path.suffix.lower() in IMAGE_EXTS:
        return 1
    if _has("fitz"):
        import fitz

        try:
            doc = fitz.open(str(path))
            n = doc.page_count
            doc.close()
            return n
        except Exception:
            return 1
    return 1


def rasterize(path: str | Path, dpi: int = 150) -> list:
    """Render each page to a PIL image (for the OCR channel). Returns [] if the
    imaging libs aren't available."""
    path = Path(path)
    images = []
    if path.suffix.lower() in IMAGE_EXTS:
        if _has("PIL"):
            from PIL import Image

            try:
                images.append(Image.open(str(path)).convert("RGB"))
            except Exception:
                pass
        return images
    if _has("fitz") and _has("PIL"):
        import fitz
        from PIL import Image

        try:
            doc = fitz.open(str(path))
            zoom = dpi / 72
            mat = fitz.Matrix(zoom, zoom)
            for pno in range(doc.page_count):
                pix = doc.load_page(pno).get_pixmap(matrix=mat)
                img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
                images.append(img)
            doc.close()
        except Exception:
            pass
    return images