extractor.py

"""Document processing: OpenCV image enhancement + Docling text extraction."""

import os
import re
import tempfile
from pathlib import Path

import cv2
import numpy as np

from llm import DEFAULT_MODEL, extract_fields
from llmplus import extract_fields_high

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

_PRICE_RE = re.compile(r'\b\d+\.\d{2}\b')


def _needs_image_fallback(text: str) -> bool:
    """
    Return True when Docling's fast-path likely missed table content:
    the output contains image-placeholder markers but almost no price values,
    meaning a table was silently swallowed as an embedded image.
    """
    markers = text.count("<!-- image -->")
    prices  = len(_PRICE_RE.findall(text))
    return markers >= 1 and prices < 3


def _render_pdf_page_as_image(file_path: str, page_num: int, dpi: int = 250) -> str:
    """
    Render a single PDF page to a PNG temp file using pypdfium2.
    Returns the temp file path; caller must delete it.
    """
    import pypdfium2 as pdfium
    doc  = pdfium.PdfDocument(file_path)
    page = doc[page_num - 1]            # 0-indexed
    scale   = dpi / 72                  # pypdfium2 native resolution is 72 dpi
    bitmap  = page.render(scale=scale, rotation=0)
    pil_img = bitmap.to_pil()
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
    pil_img.save(tmp.name)
    tmp.close()
    doc.close()
    return tmp.name

# Matches the AcceleratorDevice enum values in docling
DEVICE_AUTO = "auto"
DEVICE_CPU  = "cpu"
DEVICE_GPU  = "cuda"


# ── Docling pipeline builder ───────────────────────────────────────────────────

def _make_converter(full_pipeline: bool, device: str = DEVICE_AUTO):
    """
    full_pipeline=False  (standard)  — TableFormer disabled, 2-4x faster.
    full_pipeline=True   (high-acc)  — full ML stack, better table extraction.
    device               — 'auto' | 'cpu' | 'cuda'  passed to AcceleratorOptions.
    """
    from docling.document_converter import DocumentConverter, PdfFormatOption
    from docling.datamodel.pipeline_options import (
        PdfPipelineOptions,
        AcceleratorOptions,
        AcceleratorDevice,
    )
    from docling.datamodel.base_models import InputFormat

    device_enum = AcceleratorDevice(device)

    opts = PdfPipelineOptions()
    opts.do_table_structure = full_pipeline
    opts.do_ocr = True
    opts.accelerator_options = AcceleratorOptions(device=device_enum)

    return DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(pipeline_options=opts)
        }
    )


# ── OpenCV image preprocessing ────────────────────────────────────────────────

def preprocess_image(src: str) -> str:
    """
    Enhance a scanned document image for better OCR quality.
    Returns path to the processed image; caller is responsible for cleanup.
    """
    img = cv2.imread(src)
    if img is None:
        return src

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    h, w = gray.shape
    if min(h, w) < 1000:
        scale = max(1000 / min(h, w), 1.5)
        gray = cv2.resize(gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)

    blur_score = cv2.Laplacian(gray, cv2.CV_64F).var()

    if blur_score < 150:
        denoised = cv2.fastNlMeansDenoising(
            gray, h=10, templateWindowSize=7, searchWindowSize=21
        )
        processed = cv2.adaptiveThreshold(
            denoised, 255,
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,
            blockSize=15, C=4,
        )
    else:
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        processed = clahe.apply(gray)

    dst = src.rsplit(".", 1)[0] + "_cv2.png"
    cv2.imwrite(dst, processed)
    return dst


# ── Docling text extraction ───────────────────────────────────────────────────

def count_pdf_pages(file_path: str) -> int:
    """Return number of pages in a PDF (1 for images)."""
    ext = Path(file_path).suffix.lower()
    if ext not in {".pdf"}:
        return 1
    try:
        import pypdfium2 as pdfium
        return len(pdfium.PdfDocument(file_path))
    except Exception:
        return 1


def get_page_text(
    file_path: str,
    page_num: int,
    full_pipeline: bool = False,
    device: str = DEVICE_AUTO,
) -> str:
    """
    Extract markdown text from a single page (1-indexed).

    For PDFs: tries OCR-off first (digital PDF fast path ~1-2s).
    If the result is nearly empty (scanned page), re-runs with OCR.
    Images always use OCR.
    """
    ext = Path(file_path).suffix.lower()
    is_image = ext in IMAGE_EXTS

    if is_image:
        converter = _make_converter(full_pipeline=full_pipeline, device=device)
        result = converter.convert(file_path, page_range=(page_num, page_num))
        return result.document.export_to_markdown()

    # Digital PDF fast path — no OCR model
    from docling.datamodel.pipeline_options import PdfPipelineOptions, AcceleratorOptions, AcceleratorDevice
    from docling.document_converter import DocumentConverter, PdfFormatOption
    from docling.datamodel.base_models import InputFormat

    fast_opts = PdfPipelineOptions()
    fast_opts.do_ocr = False
    fast_opts.do_table_structure = full_pipeline
    fast_converter = DocumentConverter(
        format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=fast_opts)}
    )
    result = fast_converter.convert(file_path, page_range=(page_num, page_num))
    text = result.document.export_to_markdown()

    # Enough text AND no sign of swallowed tables → digital PDF is complete
    if len(text.strip()) >= 80 and not _needs_image_fallback(text):
        return text

    # Fallback 1: scanned page or OCR-required content
    ocr_converter = _make_converter(full_pipeline=full_pipeline, device=device)
    result = ocr_converter.convert(file_path, page_range=(page_num, page_num))
    text = result.document.export_to_markdown()

    # Fallback 2: Docling OCR still missed embedded image regions — render page
    # as a bitmap and re-run through the image OCR pipeline
    if _needs_image_fallback(text):
        img_path = _render_pdf_page_as_image(file_path, page_num)
        try:
            img_converter = _make_converter(full_pipeline=full_pipeline, device=device)
            result = img_converter.convert(img_path)
            text = result.document.export_to_markdown()
        finally:
            try:
                os.unlink(img_path)
            except OSError:
                pass

    return text


def extract_text(
    file_path: str,
    full_pipeline: bool = False,
    device: str = DEVICE_AUTO,
    progress_cb=None,
) -> str:
    """
    Convert a PDF or image to markdown via Docling (all pages concatenated).

    progress_cb: optional callable(current_page, total_pages) called after
                 each page completes. Used by app.py to drive a progress bar.
    """
    from docling.datamodel.document import ConversionResult

    converter = _make_converter(full_pipeline=full_pipeline, device=device)
    total = count_pdf_pages(file_path)

    if total <= 1 or progress_cb is None:
        result = converter.convert(file_path)
        if progress_cb:
            progress_cb(1, 1)
        text = result.document.export_to_markdown()

        # If Docling swallowed table regions as images, re-run on a bitmap render
        ext = Path(file_path).suffix.lower()
        if ext not in IMAGE_EXTS and _needs_image_fallback(text):
            img_path = _render_pdf_page_as_image(file_path, 1)
            try:
                img_converter = _make_converter(full_pipeline=full_pipeline, device=device)
                result = img_converter.convert(img_path)
                text = result.document.export_to_markdown()
            finally:
                try:
                    os.unlink(img_path)
                except OSError:
                    pass

        return text

    # Multi-page: convert page by page so we can report progress
    pages_text: list[str] = []
    for page_num in range(1, total + 1):
        result: ConversionResult = converter.convert(
            file_path,
            page_range=(page_num, page_num),
        )
        pages_text.append(result.document.export_to_markdown())
        if progress_cb:
            progress_cb(page_num, total)

    return "\n\n".join(pages_text)


# ── Shared preprocessing ──────────────────────────────────────────────────────

def _prepare_text(
    file_path: str,
    full_pipeline: bool = False,
    device: str = DEVICE_AUTO,
    progress_cb=None,
) -> str:
    ext = Path(file_path).suffix.lower()
    target = preprocess_image(file_path) if ext in IMAGE_EXTS else file_path
    text = extract_text(target, full_pipeline=full_pipeline, device=device, progress_cb=progress_cb)
    if target != file_path and os.path.exists(target):
        try:
            os.unlink(target)
        except OSError:
            pass
    return text


# ── Public pipeline functions ─────────────────────────────────────────────────

def process_file(
    file_path: str,
    model: str = DEFAULT_MODEL,
    device: str = DEVICE_AUTO,
    progress_cb=None,
) -> tuple[str, dict]:
    """Standard: OpenCV -> Docling (no TableFormer) -> single LLM pass."""
    text = _prepare_text(file_path, full_pipeline=False, device=device, progress_cb=progress_cb)
    data = extract_fields(text, model=model)
    return text, data


def process_file_high(
    file_path: str,
    model: str = DEFAULT_MODEL,
    device: str = DEVICE_AUTO,
    progress_cb=None,
) -> tuple[str, dict]:
    """High-accuracy: OpenCV -> Docling (full ML) -> 3 LLM passes + voting + math validation."""
    text = _prepare_text(file_path, full_pipeline=True, device=device, progress_cb=progress_cb)
    data = extract_fields_high(text, model=model, runs=3)
    return text, data


def detect_cuda() -> bool:
    """Return True if a CUDA-capable GPU is available to torch."""
    try:
        import torch
        return torch.cuda.is_available()
    except Exception:
        return False