utils/pdf_utils.py

"""
InvoiceForge AI — utils/pdf_utils.py
PDF to image conversion and multi-page data merging.

Uses:
  - pdf2image (poppler wrapper) for high-DPI page rendering
  - PyMuPDF (fitz) as a fallback if poppler is not installed

Provides:
  - PDFProcessor: convert PDF bytes/path to list of BGR numpy arrays
  - merge_page_results(): combine extraction results from all pages
"""

from __future__ import annotations

import io
import logging
import os
from pathlib import Path
from typing import Generator, Optional

import cv2
import numpy as np
from PIL import Image

logger = logging.getLogger(__name__)

DEFAULT_DPI = 300          # High-quality render for OCR
MAX_PAGES = 50             # Safety cap for very long documents


# ─────────────────────────────────────────────────────────────────────────────
# PDF PROCESSOR
# ─────────────────────────────────────────────────────────────────────────────

class PDFProcessor:
    """
    Converts PDF files to sequences of BGR images for OCR processing.

    Tries pdf2image first (requires poppler-utils system package),
    falls back to PyMuPDF (fitz) if poppler is unavailable.
    """

    @staticmethod
    def pdf_to_images(
        pdf_bytes: bytes,
        dpi: int = DEFAULT_DPI,
        max_pages: int = MAX_PAGES,
    ) -> list[np.ndarray]:
        """
        Convert PDF bytes to a list of BGR numpy arrays (one per page).

        Args:
            pdf_bytes: Raw PDF file contents.
            dpi:       Rendering resolution (300 recommended for OCR).
            max_pages: Maximum pages to process.

        Returns:
            List of BGR numpy arrays, one per page.

        Raises:
            ValueError: If the PDF cannot be decoded.
        """
        images: list[np.ndarray] = []

        # Try pdf2image (poppler)
        try:
            images = PDFProcessor._convert_with_pdf2image(
                pdf_bytes, dpi, max_pages
            )
            logger.info("PDF converted via pdf2image: %d pages.", len(images))
            return images
        except Exception as exc:
            logger.warning("pdf2image failed (%s), trying PyMuPDF …", exc)

        # Fallback: PyMuPDF
        try:
            images = PDFProcessor._convert_with_pymupdf(
                pdf_bytes, dpi, max_pages
            )
            logger.info("PDF converted via PyMuPDF: %d pages.", len(images))
            return images
        except Exception as exc:
            logger.error("PyMuPDF also failed: %s", exc)
            raise ValueError(
                "Could not render PDF. Ensure poppler-utils or PyMuPDF is installed."
            ) from exc

    @staticmethod
    def _convert_with_pdf2image(
        pdf_bytes: bytes,
        dpi: int,
        max_pages: int,
    ) -> list[np.ndarray]:
        """Use pdf2image (poppler) to render pages."""
        from pdf2image import convert_from_bytes  # type: ignore

        pil_pages: list[Image.Image] = convert_from_bytes(
            pdf_bytes,
            dpi=dpi,
            first_page=1,
            last_page=max_pages,
            fmt="JPEG",
        )

        bgr_pages: list[np.ndarray] = []
        for pil_img in pil_pages:
            rgb = np.array(pil_img.convert("RGB"))
            bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
            bgr_pages.append(bgr)
        return bgr_pages

    @staticmethod
    def _convert_with_pymupdf(
        pdf_bytes: bytes,
        dpi: int,
        max_pages: int,
    ) -> list[np.ndarray]:
        """Use PyMuPDF (fitz) to render pages."""
        import fitz  # type: ignore  # PyMuPDF

        pdf_doc = fitz.open(stream=pdf_bytes, filetype="pdf")
        bgr_pages: list[np.ndarray] = []
        zoom = dpi / 72.0  # 72 DPI is PDF default
        mat = fitz.Matrix(zoom, zoom)

        for page_num in range(min(len(pdf_doc), max_pages)):
            page = pdf_doc.load_page(page_num)
            pixmap = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB)
            img_bytes = pixmap.tobytes("jpeg")
            nparr = np.frombuffer(img_bytes, np.uint8)
            img_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            if img_bgr is not None:
                bgr_pages.append(img_bgr)

        pdf_doc.close()
        return bgr_pages

    @staticmethod
    def pdf_from_request(request) -> bytes:
        """
        Extract PDF bytes from a Flask request.
        Accepts multipart/form-data with key 'file'.

        Raises:
            ValueError: If no PDF found.
        """
        if request.files:
            file_obj = (
                request.files.get("file")
                or request.files.get("pdf")
                or request.files.get("image")
            )
            if file_obj:
                return file_obj.read()
        raise ValueError(
            "No PDF found in request. "
            "Send a PDF as 'file' in multipart/form-data."
        )


# ─────────────────────────────────────────────────────────────────────────────
# MULTI-PAGE RESULT MERGER
# ─────────────────────────────────────────────────────────────────────────────

def merge_page_results(page_results: list[dict]) -> dict:
    """
    Merge extraction results from multiple pages into a single document result.

    Strategy:
      - header: taken from the first page that has non-empty fields
      - items:  concatenated from all pages
      - totals: accumulated (subtotal, gstTotal, grandTotal are summed, but
                if any single page has a "grandTotal" that looks like the
                real document total it takes precedence)
      - validation: combined errors from all pages
      - meta: page count, per-page confidence

    Args:
        page_results: List of result dicts, one per page.

    Returns:
        Merged result dict.
    """
    if not page_results:
        return _empty_result()

    merged_header: dict = {}
    merged_items: list[dict] = []
    merged_errors: list[str] = []
    per_page_confidence: list[float] = []

    subtotal: float = 0.0
    gst_total: float = 0.0
    grand_total: float = 0.0
    doc_grand_total: Optional[float] = None

    for page_idx, result in enumerate(page_results):
        header = result.get("header", {})
        items = result.get("items", [])
        totals = result.get("totals", {})
        validation = result.get("validation", {})
        meta = result.get("meta", {})

        # ── Header: prefer the first page with real data ──────────────────
        if not merged_header or not merged_header.get("vendorName"):
            merged_header = _merge_header(merged_header, header)

        # ── Items: accumulate across pages ────────────────────────────────
        for item in items:
            item["_page"] = page_idx + 1
            merged_items.append(item)

        # ── Totals ─────────────────────────────────────────────────────────
        page_sub = float(totals.get("subtotal", 0.0))
        page_gst = float(totals.get("gstTotal", 0.0))
        page_gt = float(totals.get("grandTotal", 0.0))

        subtotal += page_sub
        gst_total += page_gst

        # Last page with a non-zero grandTotal is the document total
        if page_gt > 0:
            doc_grand_total = page_gt

        # ── Errors & confidence ────────────────────────────────────────────
        errors = validation.get("errors", [])
        merged_errors.extend(
            [f"[Page {page_idx + 1}] {e}" for e in errors]
        )
        per_page_confidence.append(float(validation.get("confidence", 0.85)))

    grand_total = doc_grand_total if doc_grand_total else round(subtotal + gst_total, 2)
    mean_conf = float(np.mean(per_page_confidence)) if per_page_confidence else 0.85

    return {
        "header": merged_header,
        "items": merged_items,
        "totals": {
            "subtotal": round(subtotal, 2),
            "gstTotal": round(gst_total, 2),
            "grandTotal": round(grand_total, 2),
        },
        "validation": {
            "isValid": len(merged_errors) == 0,
            "confidence": round(mean_conf, 4),
            "errors": merged_errors,
            "warnings": [],
        },
        "meta": {
            "documentType": page_results[0].get("meta", {}).get("documentType", "unknown"),
            "pageCount": len(page_results),
            "perPageConfidence": per_page_confidence,
        },
    }


def _merge_header(existing: dict, new_header: dict) -> dict:
    """Fill missing header fields from new_header into existing."""
    merged = dict(existing)
    for key, value in new_header.items():
        if value and not merged.get(key):
            merged[key] = value
    return merged


def _empty_result() -> dict:
    """Return an empty result structure."""
    return {
        "header": {
            "vendorName": "", "vendorAddress": "", "vendorGSTIN": "",
            "buyerName": "", "buyerGSTIN": "", "invoiceNumber": "",
            "invoiceDate": "", "poNumber": "",
        },
        "items": [],
        "totals": {"subtotal": 0.0, "gstTotal": 0.0, "grandTotal": 0.0},
        "validation": {
            "isValid": False, "confidence": 0.0,
            "errors": ["Empty PDF"], "warnings": [],
        },
        "meta": {"documentType": "unknown", "pageCount": 0},
    }