mnemosyne/mnemosyne/library/services/parsers.py

"""
Universal document parsing service using PyMuPDF.

Handles text extraction and image extraction for all supported formats:
PDF, EPUB, DOCX, PPTX, XLSX, XPS, MOBI, FB2, CBZ, TXT, HTML, and images.
"""

import logging
import os
import tempfile
from dataclasses import dataclass, field

import fitz  # PyMuPDF

from library.metrics import (
    DOCUMENT_PARSE_DURATION,
    DOCUMENTS_PARSED_TOTAL,
    IMAGES_EXTRACTED_TOTAL,
)

from .text_utils import remove_excessive_whitespace, sanitize_text

logger = logging.getLogger(__name__)


class UnsupportedFileTypeError(ValueError):
    """Raised when a file's type cannot be parsed.

    Deterministic and input-driven — re-parsing identical bytes can never
    succeed — so ingest must treat it as a terminal failure and never retry.
    Subclasses ``ValueError`` so existing ``except ValueError`` callers still
    catch it.
    """


# File extensions supported by PyMuPDF
PYMUPDF_EXTENSIONS = {
    "pdf", "epub", "xps", "mobi", "fb2", "cbz", "svg",
    "docx", "pptx", "xlsx", "hwpx",
}

# Plain text extensions — read directly, no PyMuPDF needed
PLAINTEXT_EXTENSIONS = {"txt", "md", "csv", "tsv", "log", "json", "yaml", "yml", "xml"}

# Image extensions — store as Image nodes directly.
# SVG is deliberately absent: it is vector XML that Pillow cannot decode and
# that the vision stage cannot send as a data URI, so it gets rasterized by
# _parse_svg_file instead.
IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "gif", "bmp", "tiff", "tif", "webp"}

# Minimum image dimensions to extract (skip tiny icons/bullets)
MIN_IMAGE_WIDTH = 50
MIN_IMAGE_HEIGHT = 50


@dataclass
class TextBlock:
    """A block of extracted text with page/section context."""

    text: str
    page: int
    metadata: dict = field(default_factory=dict)


@dataclass
class ExtractedImage:
    """An image extracted from a document."""

    data: bytes
    ext: str
    width: int
    height: int
    source_page: int
    source_index: int


@dataclass
class ParseResult:
    """Result of parsing a document: text blocks + images + metadata."""

    text_blocks: list[TextBlock] = field(default_factory=list)
    images: list[ExtractedImage] = field(default_factory=list)
    metadata: dict = field(default_factory=dict)
    file_type: str = ""


class DocumentParser:
    """
    Universal document parser using PyMuPDF.

    Extracts text and images from all supported document formats through
    a single unified interface.
    """

    def parse(self, file_path: str, file_type: str) -> ParseResult:
        """
        Parse a document and extract text blocks and images.

        :param file_path: Path to the document file.
        :param file_type: File extension (without dot), e.g. 'pdf', 'epub'.
        :returns: ParseResult with text blocks, images, and metadata.
        :raises UnsupportedFileTypeError: If the file type is not supported.
        """
        file_type = file_type.lower().lstrip(".")

        logger.info(
            "Parsing document file_type=%s path=%s",
            file_type,
            os.path.basename(file_path),
        )

        if file_type in PLAINTEXT_EXTENSIONS:
            return self._parse_plaintext(file_path, file_type)

        # Checked before PYMUPDF_EXTENSIONS: PyMuPDF can open an SVG, but
        # rendering a multi-page Write note as one document yields a blank or
        # illegible image (see svg_raster), so it needs page-aware handling.
        if file_type == "svg":
            return self._parse_svg_file(file_path, file_type)

        if file_type in IMAGE_EXTENSIONS:
            return self._parse_image_file(file_path, file_type)

        if file_type in PYMUPDF_EXTENSIONS:
            return self._parse_with_pymupdf(file_path, file_type)

        # HTML can be handled by PyMuPDF or direct read
        if file_type in ("html", "htm"):
            return self._parse_with_pymupdf(file_path, file_type)

        raise UnsupportedFileTypeError(
            f"Unsupported file type '{file_type}'. "
            f"Supported: {sorted(PYMUPDF_EXTENSIONS | PLAINTEXT_EXTENSIONS | IMAGE_EXTENSIONS)}"
        )

    def parse_bytes(self, data: bytes, file_type: str, filename: str = "") -> ParseResult:
        """
        Parse document from bytes (e.g. from S3 download).

        :param data: Raw file bytes.
        :param file_type: File extension (without dot).
        :param filename: Optional original filename for logging.
        :returns: ParseResult.
        """
        suffix = f".{file_type.lower().lstrip('.')}"
        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
            tmp.write(data)
            tmp_path = tmp.name

        try:
            logger.debug(
                "Parsing from bytes file_type=%s size=%d filename=%s",
                file_type,
                len(data),
                filename,
            )
            return self.parse(tmp_path, file_type)
        finally:
            os.unlink(tmp_path)

    def _parse_with_pymupdf(self, file_path: str, file_type: str) -> ParseResult:
        """
        Parse a document using PyMuPDF for text and image extraction.

        :param file_path: Path to the document.
        :param file_type: Normalized file extension.
        :returns: ParseResult.
        """
        with DOCUMENT_PARSE_DURATION.labels(file_type=file_type).time():
            try:
                doc = fitz.open(file_path)
            except Exception as exc:
                DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="error").inc()
                logger.error("Failed to open document file_type=%s: %s", file_type, exc)
                raise

            text_blocks: list[TextBlock] = []
            images: list[ExtractedImage] = []
            image_global_index = 0

            for page_num in range(len(doc)):
                page = doc[page_num]

                # --- Text extraction ---
                try:
                    text = page.get_text("text")
                    if text and text.strip():
                        cleaned = sanitize_text(text, log_changes=False)
                        cleaned = remove_excessive_whitespace(cleaned)
                        if cleaned.strip():
                            text_blocks.append(
                                TextBlock(text=cleaned, page=page_num)
                            )
                            logger.debug(
                                "Extracted text page=%d chars=%d",
                                page_num,
                                len(cleaned),
                            )
                except Exception as exc:
                    logger.warning(
                        "Text extraction failed page=%d: %s, continuing",
                        page_num,
                        exc,
                    )

                # --- Image extraction ---
                try:
                    for img_info in page.get_images(full=True):
                        xref = img_info[0]
                        try:
                            img_data = doc.extract_image(xref)
                            if not img_data or not img_data.get("image"):
                                continue

                            width = img_data.get("width", 0)
                            height = img_data.get("height", 0)

                            # Skip tiny images (icons, bullets, etc.)
                            if width < MIN_IMAGE_WIDTH or height < MIN_IMAGE_HEIGHT:
                                logger.debug(
                                    "Skipping small image page=%d xref=%d size=%dx%d",
                                    page_num,
                                    xref,
                                    width,
                                    height,
                                )
                                continue

                            images.append(
                                ExtractedImage(
                                    data=img_data["image"],
                                    ext=img_data.get("ext", "png"),
                                    width=width,
                                    height=height,
                                    source_page=page_num,
                                    source_index=image_global_index,
                                )
                            )
                            image_global_index += 1

                            logger.debug(
                                "Extracted image page=%d format=%s size=%dx%d bytes=%d",
                                page_num,
                                img_data.get("ext", "?"),
                                width,
                                height,
                                len(img_data["image"]),
                            )
                        except Exception as exc:
                            logger.warning(
                                "Image extraction failed page=%d xref=%d: %s",
                                page_num,
                                xref,
                                exc,
                            )
                except Exception as exc:
                    logger.warning(
                        "Image listing failed page=%d: %s, continuing",
                        page_num,
                        exc,
                    )

            # Collect document metadata
            meta = doc.metadata or {}
            result_meta = {
                "page_count": len(doc),
                "title": meta.get("title", ""),
                "author": meta.get("author", ""),
                "subject": meta.get("subject", ""),
                "creator": meta.get("creator", ""),
            }

            doc.close()

            DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="success").inc()
            IMAGES_EXTRACTED_TOTAL.labels(file_type=file_type).inc(len(images))

            logger.info(
                "Parsed document file_type=%s pages=%d text_blocks=%d images=%d",
                file_type,
                result_meta["page_count"],
                len(text_blocks),
                len(images),
            )

            return ParseResult(
                text_blocks=text_blocks,
                images=images,
                metadata=result_meta,
                file_type=file_type,
            )

    def _parse_plaintext(self, file_path: str, file_type: str) -> ParseResult:
        """
        Parse a plain text file by direct read.

        :param file_path: Path to the text file.
        :param file_type: Normalized file extension.
        :returns: ParseResult.
        """
        with DOCUMENT_PARSE_DURATION.labels(file_type=file_type).time():
            try:
                with open(file_path, "r", encoding="utf-8", errors="replace") as f:
                    content = f.read()
            except Exception as exc:
                DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="error").inc()
                logger.error("Failed to read text file file_type=%s: %s", file_type, exc)
                raise

            cleaned = sanitize_text(content, log_changes=True)
            cleaned = remove_excessive_whitespace(cleaned)

            text_blocks = []
            if cleaned.strip():
                text_blocks.append(TextBlock(text=cleaned, page=0))

            DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="success").inc()

            logger.info(
                "Parsed plaintext file_type=%s chars=%d",
                file_type,
                len(cleaned),
            )

            return ParseResult(
                text_blocks=text_blocks,
                images=[],
                metadata={"page_count": 1},
                file_type=file_type,
            )

    def _parse_svg_file(self, file_path: str, file_type: str) -> ParseResult:
        """
        Rasterize an SVG into one ExtractedImage per page.

        SVG is vector XML: Pillow cannot decode it and the vision stage cannot
        put it in a data URI, so it is rendered to PNG here. Multi-page Write
        notes become one image per page so each page reaches vision/OCR at a
        legible size.

        :param file_path: Path to the SVG file.
        :param file_type: Normalized file extension ("svg").
        :returns: ParseResult with one image per rendered page.
        """
        with DOCUMENT_PARSE_DURATION.labels(file_type=file_type).time():
            try:
                from library.services.svg_raster import render_svg_pages

                with open(file_path, "rb") as f:
                    data = f.read()

                pages = render_svg_pages(data)
            except Exception as exc:
                DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="error").inc()
                logger.error("Failed to rasterize SVG file_type=%s: %s", file_type, exc)
                raise

            images = [
                ExtractedImage(
                    data=png,
                    ext="png",
                    width=width,
                    height=height,
                    source_page=index,
                    source_index=0,
                )
                for index, (png, width, height) in enumerate(pages)
            ]

            DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="success").inc()
            IMAGES_EXTRACTED_TOTAL.labels(file_type=file_type).inc(len(images))

            logger.info(
                "Parsed SVG file_type=%s pages=%d bytes=%d",
                file_type,
                len(images),
                len(data),
            )

            return ParseResult(
                text_blocks=[],
                images=images,
                metadata={"page_count": len(images)},
                file_type=file_type,
            )

    def _parse_image_file(self, file_path: str, file_type: str) -> ParseResult:
        """
        Handle a standalone image file — store as a single ExtractedImage.

        :param file_path: Path to the image file.
        :param file_type: Normalized file extension.
        :returns: ParseResult with one image and no text.
        """
        with DOCUMENT_PARSE_DURATION.labels(file_type=file_type).time():
            try:
                from PIL import Image as PILImage

                with open(file_path, "rb") as f:
                    data = f.read()

                img = PILImage.open(file_path)
                width, height = img.size
                img.close()
            except Exception as exc:
                DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="error").inc()
                logger.error("Failed to read image file_type=%s: %s", file_type, exc)
                raise

            DOCUMENTS_PARSED_TOTAL.labels(file_type=file_type, status="success").inc()
            IMAGES_EXTRACTED_TOTAL.labels(file_type=file_type).inc(1)

            logger.info(
                "Parsed image file file_type=%s size=%dx%d bytes=%d",
                file_type,
                width,
                height,
                len(data),
            )

            return ParseResult(
                text_blocks=[],
                images=[
                    ExtractedImage(
                        data=data,
                        ext=file_type,
                        width=width,
                        height=height,
                        source_page=0,
                        source_index=0,
                    )
                ],
                metadata={"page_count": 0, "width": width, "height": height},
                file_type=file_type,
            )