mcp-nextcloud/nextcloud_mcp_server/document_processors/pymupdf.py

"""Document processor using PyMuPDF (fitz) library."""

import logging
import pathlib
import tempfile
from collections.abc import Awaitable, Callable
from typing import Any, Optional

import anyio

# NOTE: Do NOT call pymupdf.layout.activate() here!
# It changes the behavior of pymupdf4llm.to_markdown() when page_chunks=True,
# causing it to return a string instead of a list[dict].
# See: https://github.com/pymupdf/pymupdf4llm/issues/323
import pymupdf
import pymupdf4llm

from .base import DocumentProcessor, ProcessingResult, ProcessorError

logger = logging.getLogger(__name__)


class PyMuPDFProcessor(DocumentProcessor):
    """Document processor using PyMuPDF library for PDF processing.

    PyMuPDF (fitz) is a fast, local PDF processing library that extracts text,
    metadata, and images without requiring external API calls.

    Features:
    - Fast text extraction with layout preservation
    - PDF metadata extraction (title, author, creation date, page count)
    - Image extraction for future multimodal support
    - Page number tracking for precise citations
    """

    SUPPORTED_TYPES = {
        "application/pdf",
    }

    def __init__(
        self,
        extract_images: bool = True,
        image_dir: Optional[str | pathlib.Path] = None,
    ):
        """Initialize PyMuPDF processor.

        Args:
            extract_images: Whether to extract embedded images from PDFs
            image_dir: Directory to store extracted images (defaults to temp directory)
        """
        self.extract_images = extract_images

        if image_dir is None:
            self.image_dir = pathlib.Path(tempfile.gettempdir()) / "pdf-images"
        else:
            self.image_dir = pathlib.Path(image_dir)

        # Create image directory if it doesn't exist
        if self.extract_images:
            self.image_dir.mkdir(exist_ok=True, parents=True)
            logger.info(
                "Initialized PyMuPDFProcessor with image extraction to %s",
                self.image_dir,
            )
        else:
            logger.info("Initialized PyMuPDFProcessor without image extraction")

    @property
    def name(self) -> str:
        return "pymupdf"

    @property
    def supported_mime_types(self) -> set[str]:
        return self.SUPPORTED_TYPES

    async def process(
        self,
        content: bytes,
        content_type: str,
        filename: Optional[str] = None,
        options: Optional[dict[str, Any]] = None,
        progress_callback: Optional[
            Callable[[float, Optional[float], Optional[str]], Awaitable[None]]
        ] = None,
    ) -> ProcessingResult:
        """Process a PDF document and extract text, metadata, and images.

        Args:
            content: PDF document bytes
            content_type: MIME type (should be application/pdf)
            filename: Optional filename for better error messages
            options: Processing options (currently unused)
            progress_callback: Optional callback for progress updates

        Returns:
            ProcessingResult with extracted text and metadata

        Raises:
            ProcessorError: If PDF processing fails
        """

        try:
            if progress_callback:
                await progress_callback(0, 100, "Opening PDF document")

            # Open document and extract metadata in thread
            doc = await anyio.to_thread.run_sync(  # type: ignore[attr-defined]
                lambda: pymupdf.open("pdf", content)
            )

            metadata = self._extract_metadata(doc, filename)
            metadata["file_size"] = len(content)
            page_count = doc.page_count

            if progress_callback:
                await progress_callback(10, 100, f"Extracting {page_count} pages")

            # Prepare image directory if needed
            pdf_image_dir = None
            if self.extract_images:
                pdf_id = filename.replace("/", "_") if filename else "unknown"
                pdf_image_dir = self.image_dir / pdf_id
                pdf_image_dir.mkdir(exist_ok=True, parents=True)

            # Extract all pages in a single call with page_chunks=True
            def do_extract() -> list[dict[str, Any]]:
                # When page_chunks=True, to_markdown returns list[dict] not str
                return pymupdf4llm.to_markdown(  # type: ignore[return-value]
                    doc,
                    write_images=self.extract_images,
                    image_path=pdf_image_dir if self.extract_images else None,
                    page_chunks=True,
                )

            page_chunks: list[dict[str, Any]] = await anyio.to_thread.run_sync(  # type: ignore[attr-defined]
                do_extract
            )

            if progress_callback:
                await progress_callback(90, 100, "Building result")

            # Extract page texts and build boundaries from chunks
            page_texts: list[str] = []
            page_boundaries: list[dict[str, Any]] = []
            current_offset = 0
            for chunk in page_chunks:
                text = chunk.get("text", "")
                page_num = chunk.get("metadata", {}).get("page", len(page_texts) + 1)
                page_texts.append(text)
                page_boundaries.append(
                    {
                        "page": page_num,
                        "start_offset": current_offset,
                        "end_offset": current_offset + len(text),
                    }
                )
                current_offset += len(text)

            # Collect image paths
            image_paths = []
            if pdf_image_dir and pdf_image_dir.exists():
                image_paths = [str(p) for p in pdf_image_dir.glob("*")]

            # Build final text and metadata
            md_text = "".join(page_texts)
            metadata["has_images"] = len(image_paths) > 0
            if image_paths:
                metadata["image_count"] = len(image_paths)
                metadata["image_paths"] = image_paths
            metadata["page_boundaries"] = page_boundaries

            # Close document
            doc.close()

            if progress_callback:
                await progress_callback(100, 100, "Processing complete")

            logger.info(
                "Successfully processed PDF %s: %s pages, %s chars, %s images",
                filename or "<bytes>",
                metadata["page_count"],
                len(md_text),
                metadata.get("image_count", 0),
                extra={
                    "processor": self.name,
                    "tier": self.tier,
                    "pages": metadata["page_count"],
                    "chars": len(md_text),
                    "images": metadata.get("image_count", 0),
                    "byte_size": len(content),
                },
            )

            return ProcessingResult(
                text=md_text,
                metadata=metadata,
                processor=self.name,
                success=True,
            )

        except Exception as e:
            error_msg = f"Failed to process PDF {filename or '<bytes>'}: {e}"
            logger.error(error_msg, exc_info=True)
            raise ProcessorError(error_msg) from e

    def _extract_metadata(
        self, doc: pymupdf.Document, filename: Optional[str]
    ) -> dict[str, Any]:
        """Extract metadata from PDF document.

        Args:
            doc: Opened PyMuPDF document
            filename: Optional filename

        Returns:
            Dictionary with PDF metadata
        """
        metadata: dict[str, Any] = {}

        # Basic document info
        metadata["page_count"] = doc.page_count
        metadata["format"] = "PDF 1." + str(
            doc.pdf_version() if hasattr(doc, "pdf_version") else "?"  # type: ignore[call-non-callable]
        )

        if filename:
            metadata["filename"] = filename

        # Extract PDF metadata dictionary
        pdf_metadata = doc.metadata
        if pdf_metadata:
            # Standard PDF metadata fields
            if pdf_metadata.get("title"):
                metadata["title"] = pdf_metadata["title"]
            if pdf_metadata.get("author"):
                metadata["author"] = pdf_metadata["author"]
            if pdf_metadata.get("subject"):
                metadata["subject"] = pdf_metadata["subject"]
            if pdf_metadata.get("keywords"):
                metadata["keywords"] = pdf_metadata["keywords"]
            if pdf_metadata.get("creator"):
                metadata["creator"] = pdf_metadata["creator"]
            if pdf_metadata.get("producer"):
                metadata["producer"] = pdf_metadata["producer"]
            if pdf_metadata.get("creationDate"):
                metadata["creation_date"] = pdf_metadata["creationDate"]
            if pdf_metadata.get("modDate"):
                metadata["modification_date"] = pdf_metadata["modDate"]

        return metadata

    async def health_check(self) -> bool:
        """Check if PyMuPDF is available and working.

        Returns:
            True if processor is ready to use
        """
        try:
            # Try to create a simple PDF in memory
            test_doc = pymupdf.open()
            test_doc.close()
            return True
        except Exception as e:
            logger.error("PyMuPDF health check failed: %s", e)
            return False