pdf/gateway/app/services/convert/layout/pdf_router.py

"""Document-level PDF routing (pdf-inspector-inspired, in-house).

Classifies TextBased / Scanned / ImageBased / Mixed and lists pages needing OCR.
Does not call Firecrawl or pdf-inspector — pure heuristics over pypdf + optional engine.
"""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from enum import Enum

from pypdf import PdfReader

from app.services.convert import doc_cache


class PdfDocType(str, Enum):
    text_based = "text_based"
    scanned = "scanned"
    image_based = "image_based"
    mixed = "mixed"


class ScanStrategy(str, Enum):
    early_exit = "early_exit"
    full = "full"
    sample = "sample"


@dataclass
class PdfRouteResult:
    doc_type: PdfDocType
    confidence: float
    pages_needing_ocr: list[int] = field(default_factory=list)
    pages_sampled: list[int] = field(default_factory=list)
    strategy: str = "full"

    def as_dict(self) -> dict:
        return {
            "doc_type": self.doc_type.value,
            "confidence": round(self.confidence, 4),
            "pages_needing_ocr": list(self.pages_needing_ocr),
            "pages_sampled": list(self.pages_sampled),
            "strategy": self.strategy,
        }


def _strategy_from_env() -> tuple[ScanStrategy, int]:
    raw = (os.environ.get("CONVERT_PDF_SCAN_STRATEGY") or "full").strip().lower()
    try:
        n = int(os.environ.get("CONVERT_PDF_SCAN_SAMPLE", "5"))
    except ValueError:
        n = 5
    n = max(2, min(n, 50))
    if raw in ("early_exit", "early-exit", "early"):
        return ScanStrategy.early_exit, n
    if raw in ("sample", "sampled"):
        return ScanStrategy.sample, n
    return ScanStrategy.full, n


def _page_needs_ocr(reader: PdfReader, index: int, data: bytes | None = None) -> tuple[bool, str]:
    """Heuristic: sparse / garbled extractable text → OCR. Returns (needs_ocr, reason)."""
    from app.services.convert.layout.text_quality import is_garbled_text, page_text_needs_ocr

    try:
        page = reader.pages[index]
        text = (
            doc_cache.page_text(data, index) if data is not None else (page.extract_text() or "")
        ).strip()
    except Exception:
        return True, "extract_failed"

    # Image XObjects hint
    has_image = False
    try:
        resources = page.get("/Resources") or {}
        if hasattr(resources, "get_object"):
            resources = resources.get_object()
        xobj = resources.get("/XObject") if resources else None
        if xobj is not None:
            if hasattr(xobj, "get_object"):
                xobj = xobj.get_object()
            for _name, ref in (xobj.items() if hasattr(xobj, "items") else []):
                try:
                    obj = ref.get_object() if hasattr(ref, "get_object") else ref
                    if obj.get("/Subtype") == "/Image":
                        has_image = True
                        break
                except Exception:
                    continue
    except Exception:
        pass

    chars = len(text)
    # Image-only / nearly empty → OCR. Decorative images beside real text → keep digital.
    if chars < 12:
        return True, "image_only" if has_image else "sparse_text"
    # Broken ToUnicode / CID garbage (Huwiyati-class) must not count as "text_ok"
    if is_garbled_text(text) or page_text_needs_ocr(text, image_coverage=0.2 if has_image else 0.0):
        return True, "garbled_text"
    return False, "text_ok"


def _sample_indices(n_pages: int, strategy: ScanStrategy, sample_n: int) -> list[int]:
    if n_pages <= 0:
        return []
    if strategy == ScanStrategy.full or strategy == ScanStrategy.early_exit:
        return list(range(n_pages))
    # sample: first, last, evenly spaced middle
    if n_pages <= sample_n:
        return list(range(n_pages))
    idxs = {0, n_pages - 1}
    mid_slots = sample_n - 2
    if mid_slots > 0:
        step = (n_pages - 1) / (mid_slots + 1)
        for i in range(1, mid_slots + 1):
            idxs.add(min(n_pages - 1, max(0, int(round(i * step)))))
    return sorted(idxs)


def route_pdf(data: bytes, *, strategy: ScanStrategy | None = None) -> PdfRouteResult:
    """Classify PDF and list 0-based page indices that need OCR."""
    strat, sample_n = _strategy_from_env()
    if strategy is not None:
        strat = strategy

    try:
        reader = doc_cache.get_reader(data)
    except Exception:
        return PdfRouteResult(
            doc_type=PdfDocType.scanned,
            confidence=0.3,
            pages_needing_ocr=[],
            pages_sampled=[],
            strategy=strat.value,
        )

    n_pages = len(reader.pages)
    indices = _sample_indices(n_pages, strat, sample_n)
    needing: list[int] = []
    sampled: list[int] = []
    text_ok = 0
    image_heavy = 0
    garbled_heavy = 0

    for i in indices:
        sampled.append(i)
        needs, reason = _page_needs_ocr(reader, i, data)
        if needs:
            needing.append(i)
            if reason in ("image_only", "hybrid_sparse"):
                image_heavy += 1
            if reason == "garbled_text":
                garbled_heavy += 1
        else:
            text_ok += 1
        if strat == ScanStrategy.early_exit and needing and not text_ok:
            # First page already needs OCR and no prior text → treat as scanned early
            break

    # Expand early_exit / sample results to full page lists when mixed
    if strat == ScanStrategy.early_exit and needing and text_ok == 0 and len(sampled) < n_pages:
        # Likely fully scanned — mark all pages
        needing = list(range(n_pages))
        sampled = list(range(n_pages))
    elif strat == ScanStrategy.sample and needing:
        # Conservative: any sampled OCR page → include that page; do not invent others
        pass

    total = max(len(sampled), 1)
    ocr_frac = len(needing) / total
    text_frac = text_ok / total
    garbled_frac = garbled_heavy / total

    # Majority garbled / OCR-needed → promote whole document (hybrid RFP class)
    if ocr_frac >= 0.5 and len(sampled) < n_pages:
        needing = list(range(n_pages))
        sampled = list(range(n_pages))
        ocr_frac = 1.0
        text_frac = 0.0
    elif ocr_frac >= 0.5 and n_pages > 0 and len(needing) < n_pages:
        needing = list(range(n_pages))
        ocr_frac = 1.0
        text_frac = 0.0

    if ocr_frac >= 0.95:
        if garbled_frac >= 0.4:
            doc_type = PdfDocType.mixed if text_ok > 0 else PdfDocType.scanned
        else:
            doc_type = PdfDocType.image_based if image_heavy >= text_ok else PdfDocType.scanned
        conf = 0.55 + 0.4 * ocr_frac
    elif ocr_frac <= 0.05:
        doc_type = PdfDocType.text_based
        conf = 0.55 + 0.4 * text_frac
    else:
        doc_type = PdfDocType.mixed
        conf = 0.5 + 0.3 * (1.0 - abs(0.5 - ocr_frac))

    # Unique sorted
    needing = sorted(set(needing))
    return PdfRouteResult(
        doc_type=doc_type,
        confidence=min(1.0, conf),
        pages_needing_ocr=needing,
        pages_sampled=sampled,
        strategy=strat.value,
    )


def page_kind_hint(doc_type: PdfDocType, page_index: int, pages_needing_ocr: list[int]):
    """Map route into PageKind-compatible hint for a page."""
    from app.services.convert.idm.model import PageKind

    if page_index in pages_needing_ocr:
        if doc_type == PdfDocType.mixed:
            return PageKind.hybrid
        return PageKind.scan
    if doc_type == PdfDocType.text_based:
        return PageKind.digital
    if doc_type == PdfDocType.mixed:
        return PageKind.digital
    return PageKind.scan