src/extractors/document_loader.py

"""Load documents (PDF / PNG / JPG) into text + page images for extraction.

Strategy:
1. For PDFs, try text extraction first via pdfplumber. If the doc has meaningful
   text, we use it directly (fast, cheap).
2. Always also render page images via PyMuPDF — GPT-5 nano vision handles layout
   information that pure text misses (tables, spatial context on receipts).
3. Standalone images (PNG/JPG) skip text extraction and go straight to vision.

The extractor decides how to use text vs images based on `source_type`.
"""
from __future__ import annotations

import base64
import io
from dataclasses import dataclass, field
from pathlib import Path

import fitz  # PyMuPDF
import pdfplumber
from PIL import Image

from src.utils.logging import logger

# --- Config knobs ------------------------------------------------------------

# If extracted text is shorter than this, treat the PDF as scanned/image-based.
_MIN_TEXT_CHARS_FOR_TEXT_PDF = 100

# Max pages we render as images (guards against absurdly long PDFs pre-v2).
_MAX_PAGES_TO_RENDER = 5

# DPI for rendering — higher = clearer OCR but slower + bigger payload.
# 200 DPI is a good sweet spot for receipts + invoices.
_RENDER_DPI = 200

# Max side length before we downscale — GPT-5 nano vision handles up to 2048.
_MAX_IMAGE_SIDE = 2048


@dataclass
class LoadedDocument:
    """The output of the loader — text and/or images ready for the LLM."""

    text: str = ""
    images_b64: list[str] = field(default_factory=list)
    source_type: str = "unknown"  # "text_pdf" | "image_pdf" | "image" | "empty"
    page_count: int = 0
    filename: str = ""


# --- Helpers ---------------------------------------------------------------


def _image_to_b64(img: Image.Image) -> str:
    """PIL Image -> base64-encoded PNG string suitable for OpenAI vision."""
    # Downscale if needed
    if max(img.size) > _MAX_IMAGE_SIDE:
        ratio = _MAX_IMAGE_SIDE / max(img.size)
        new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio))
        img = img.resize(new_size, Image.LANCZOS)

    buf = io.BytesIO()
    img.convert("RGB").save(buf, format="PNG", optimize=True)
    return base64.b64encode(buf.getvalue()).decode("ascii")


def _render_pdf_pages(pdf_bytes: bytes, max_pages: int = _MAX_PAGES_TO_RENDER) -> list[str]:
    """Render each PDF page to a base64 PNG using PyMuPDF."""
    images_b64: list[str] = []
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        for page_num in range(min(len(doc), max_pages)):
            page = doc[page_num]
            pix = page.get_pixmap(dpi=_RENDER_DPI)
            img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
            images_b64.append(_image_to_b64(img))
    return images_b64


def _extract_pdf_text(pdf_bytes: bytes) -> str:
    """Extract text from all pages via pdfplumber. Empty string on failure."""
    try:
        with pdfplumber.open(io.BytesIO(pdf_bytes)) as pdf:
            parts: list[str] = []
            for page in pdf.pages[:_MAX_PAGES_TO_RENDER]:
                text = page.extract_text() or ""
                if text.strip():
                    parts.append(text)
            return "\n\n".join(parts).strip()
    except Exception as e:
        logger.warning(f"pdfplumber text extraction failed: {e}")
        return ""


# --- Public API ------------------------------------------------------------


def load_document(
    file_bytes: bytes,
    filename: str = "document",
    *,
    render_images: bool = True,
) -> LoadedDocument:
    """Turn raw file bytes into a LoadedDocument.

    - `render_images=False` skips image rendering for cost savings on text-heavy PDFs.
      The extractor can decide based on doc size / cost budget.
    """
    ext = Path(filename).suffix.lower()

    # --- Standalone images -> straight to vision -----------------------------
    if ext in {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff", ".tif"}:
        try:
            img = Image.open(io.BytesIO(file_bytes))
            return LoadedDocument(
                text="",
                images_b64=[_image_to_b64(img)],
                source_type="image",
                page_count=1,
                filename=filename,
            )
        except Exception as e:
            logger.error(f"Failed to load image {filename}: {e}")
            return LoadedDocument(source_type="empty", filename=filename)

    # --- PDFs ---------------------------------------------------------------
    if ext == ".pdf" or file_bytes[:4] == b"%PDF":
        text = _extract_pdf_text(file_bytes)
        has_text = len(text) >= _MIN_TEXT_CHARS_FOR_TEXT_PDF

        images_b64: list[str] = []
        if render_images or not has_text:
            try:
                images_b64 = _render_pdf_pages(file_bytes)
            except Exception as e:
                logger.warning(f"PDF image rendering failed: {e}")

        source_type = "text_pdf" if has_text else "image_pdf"

        # Get page count for logging
        try:
            with fitz.open(stream=file_bytes, filetype="pdf") as doc:
                page_count = len(doc)
        except Exception:
            page_count = len(images_b64)

        logger.info(
            f"Loaded {filename}: source_type={source_type}, pages={page_count}, "
            f"text_chars={len(text)}, images={len(images_b64)}"
        )
        return LoadedDocument(
            text=text,
            images_b64=images_b64,
            source_type=source_type,
            page_count=page_count,
            filename=filename,
        )

    # --- Plain text ---------------------------------------------------------
    # `.txt` (and other text-like extensions) come from the eval CLI's inline
    # `text` field, and from OCR outputs. No images to render — the LLM works
    # on the decoded string directly. Fall back to lossy UTF-8 decoding so a
    # stray non-UTF byte doesn't kill the whole record.
    if ext in {".txt", ".text", ".md", ".log"}:
        try:
            text = file_bytes.decode("utf-8", errors="replace").strip()
        except Exception as e:
            logger.error(f"Failed to decode text file {filename}: {e}")
            return LoadedDocument(source_type="empty", filename=filename)

        source_type = "text" if text else "empty"
        logger.info(f"Loaded {filename}: source_type={source_type}, text_chars={len(text)}")
        return LoadedDocument(
            text=text,
            images_b64=[],
            source_type=source_type,
            page_count=1 if text else 0,
            filename=filename,
        )

    # --- Unknown format -----------------------------------------------------
    logger.warning(f"Unknown file extension {ext!r} for {filename}. Treating as empty.")
    return LoadedDocument(source_type="empty", filename=filename)