core/converter.py

"""
Document converter for the MarkItDown API.

Wraps Microsoft MarkItDown and RapidOCR to provide a unified conversion
interface that accepts file paths, raw byte streams, and public URLs.

Supported extensions are declared in SUPPORTED_EXTENSIONS and imported by
the API layer for format listing and validation.

Public classes
--------------
    ConversionResult
        Immutable dataclass holding the converted Markdown and file statistics.

    ConversionError
        Immutable dataclass holding error details when conversion fails.

    DocumentConverter
        Main converter class. All convert_* methods return either a
        ConversionResult or a ConversionError — they do not raise.
"""

from __future__ import annotations

import hashlib
import io
import mimetypes
import time
from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterator, Optional
from urllib.parse import urlparse

from markitdown import MarkItDown

from logger import get_logger

from .ocr_engine import ocr_image, ocr_pdf

logger = get_logger(__name__)

IMAGE_EXTENSIONS: frozenset[str] = frozenset({
    ".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp", ".tiff",
})

IMAGE_MIME_PREFIXES: tuple[str, ...] = ("image/",)

SUPPORTED_EXTENSIONS: frozenset[str] = frozenset({
    ".pdf", ".docx", ".doc", ".pptx", ".ppt",
    ".xlsx", ".xls", ".csv", ".json", ".xml",
    ".html", ".htm", ".txt", ".md", ".rst",
    ".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp", ".tiff",
    ".mp3", ".wav", ".ogg", ".flac",
    ".zip", ".epub",
})


def _is_image(ext: str, mime: str) -> bool:
    """Return True when the input should be routed through RapidOCR."""
    return ext.lower() in IMAGE_EXTENSIONS or mime.startswith(IMAGE_MIME_PREFIXES)


# ---------------------------------------------------------------------------
# Result and error types
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class ConversionResult:
    """Successful conversion output."""

    source: str
    markdown: str
    char_count: int
    word_count: int
    line_count: int
    duration_ms: float
    file_size_bytes: int
    mime_type: str
    content_hash: str
    metadata: dict = field(default_factory=dict)

    @property
    def token_estimate(self) -> int:
        """Rough LLM token estimate based on word count (4/3 words per token)."""
        return max(1, self.word_count * 4 // 3)


@dataclass(frozen=True)
class ConversionError:
    """Conversion failure details."""

    source: str
    error_type: str
    message: str
    duration_ms: float


# ---------------------------------------------------------------------------
# Converter
# ---------------------------------------------------------------------------

class DocumentConverter:
    """Converts documents from various formats to Markdown.

    All public methods return either a ConversionResult or a ConversionError
    and never raise exceptions to callers.
    """

    def __init__(self, enable_plugins: bool = False) -> None:
        self._engine = MarkItDown(enable_plugins=enable_plugins)

    # ------------------------------------------------------------------
    # Public conversion methods
    # ------------------------------------------------------------------

    def convert_file(self, path: str | Path) -> ConversionResult | ConversionError:
        """Convert a local file identified by *path*."""
        path = Path(path).resolve()

        if not path.exists():
            return self._error(str(path), "FileNotFoundError",
                               f"File does not exist: {path}", 0.0)

        file_size = path.stat().st_size
        mime_type = mimetypes.guess_type(str(path))[0] or "application/octet-stream"
        ext = path.suffix

        with self._timed() as elapsed:
            try:
                markdown = self._convert_to_markdown(ext, mime_type, str(path))
            except Exception as exc:
                logger.error("convert_file | exception | file=%s | error=%s",
                             path, exc, exc_info=True)
                return self._error(str(path), type(exc).__name__, str(exc), elapsed())

        return self._build_result(str(path), markdown, file_size, mime_type, elapsed())

    def convert_url(self, url: str) -> ConversionResult | ConversionError:
        """Fetch and convert a public HTTP/HTTPS URL."""
        parsed = urlparse(url)
        if parsed.scheme not in {"http", "https"}:
            return self._error(url, "ValueError",
                               f"Unsupported URL scheme: {parsed.scheme!r}", 0.0)

        ext = Path(parsed.path).suffix.lower()
        is_image = ext in IMAGE_EXTENSIONS
        mime_type = mimetypes.guess_type(url)[0] or ("image/jpeg" if is_image else "text/html")

        with self._timed() as elapsed:
            try:
                markdown = self._convert_to_markdown(ext, mime_type, url)
            except Exception as exc:
                logger.error("convert_url | exception | url=%s | error=%s",
                             url, exc, exc_info=True)
                return self._error(url, type(exc).__name__, str(exc), elapsed())

        return self._build_result(url, markdown, 0, mime_type, elapsed())

    def convert_stream(self, data: bytes, filename: str) -> ConversionResult | ConversionError:
        """Convert raw bytes identified by *filename* (used for upload payloads)."""
        mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
        ext = Path(filename).suffix.lower()

        with self._timed() as elapsed:
            try:
                markdown = self._convert_to_markdown(ext, mime_type, data)
            except Exception as exc:
                logger.error("convert_stream | exception | filename=%s | error=%s",
                             filename, exc, exc_info=True)
                return self._error(filename, type(exc).__name__, str(exc), elapsed())

        return self._build_result(filename, markdown, len(data), mime_type, elapsed())

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _convert_to_markdown(self, ext: str, mime_type: str, source) -> str:
        """Dispatch to the right backend (OCR or MarkItDown) and return text.

        For PDFs, if MarkItDown yields nothing, fall back to OCR.
        `source` is a path/URL for image inputs, or raw bytes for streams.
        """
        if _is_image(ext, mime_type):
            return ocr_image(source)

        if isinstance(source, bytes):
            markdown = self._engine.convert_stream(
                io.BytesIO(source), file_extension=ext
            ).text_content
        else:
            markdown = self._engine.convert(source).text_content

        if not markdown.strip() and ext.lower() == ".pdf":
            source_label = f"<{len(source)} bytes>" if isinstance(source, (bytes, bytearray)) else source
            logger.info("convert | PDF text empty, falling back to OCR | source=%s", source_label)
            markdown = ocr_pdf(source)
        return markdown

    @staticmethod
    def _build_result(
        source: str,
        markdown: str,
        file_size: int,
        mime_type: str,
        elapsed: float,
    ) -> ConversionResult:
        return ConversionResult(
            source=source,
            markdown=markdown,
            char_count=len(markdown),
            word_count=len(markdown.split()),
            line_count=len(markdown.splitlines()),
            duration_ms=elapsed,
            file_size_bytes=file_size,
            mime_type=mime_type,
            content_hash=hashlib.sha256(markdown.encode()).hexdigest(),
        )

    @staticmethod
    def _error(source: str, error_type: str, message: str, duration_ms: float) -> ConversionError:
        return ConversionError(
            source=source,
            error_type=error_type,
            message=message,
            duration_ms=duration_ms,
        )

    @staticmethod
    @contextmanager
    def _timed() -> Iterator:
        """Context manager that yields a callable returning elapsed ms at any point."""
        start = time.perf_counter()

        def elapsed() -> float:
            return (time.perf_counter() - start) * 1000.0

        yield elapsed