src/document_parser.py

"""Docling document parsing with figure extraction and markdown export."""

from __future__ import annotations

import os
import uuid
from collections.abc import Callable
from pathlib import Path
from typing import Any

from storage import get_temp_dir


_EXT_TO_INPUT_FORMAT = {
    ".pdf": "PDF",
    ".docx": "DOCX",
    ".xlsx": "XLSX",
    ".pptx": "PPTX",
}

_converter: Any = None


def _is_zerogpu_space() -> bool:
    """Detect HuggingFace ZeroGPU Spaces where CUDA init must happen inside @spaces.GPU."""
    return bool(os.environ.get("SPACE_ID"))


def _pick_accelerator_device() -> Any:
    """Return the best AcceleratorDevice for Docling, or None if unsupported."""
    try:
        from docling.datamodel.pipeline_options import AcceleratorDevice
    except ImportError:
        return None

    if _is_zerogpu_space():
        return AcceleratorDevice.CPU

    try:
        import torch
        if torch.cuda.is_available():
            return AcceleratorDevice.CUDA
        if torch.backends.mps.is_available():
            return AcceleratorDevice.MPS
    except ImportError:
        pass

    return AcceleratorDevice.CPU


def get_converter() -> Any:
    """Return a shared DocumentConverter instance, creating it on first call."""
    global _converter  # noqa: PLW0603
    if _converter is not None:
        return _converter

    from docling.datamodel.base_models import InputFormat
    from docling.document_converter import DocumentConverter, PdfFormatOption

    pdf_format_option = PdfFormatOption()
    pdf_format_option.pipeline_options.generate_picture_images = True
    pdf_format_option.pipeline_options.images_scale = 2.0

    device = _pick_accelerator_device()
    if device is not None:
        try:
            from docling.datamodel.pipeline_options import AcceleratorOptions
            pdf_format_option.pipeline_options.accelerator_options = AcceleratorOptions(
                device=device,
            )
        except Exception:  # noqa: BLE001
            pass

    # On ZeroGPU, hide CUDA from sub-libraries (RapidOCR, etc.) that do their
    # own torch.cuda detection and trigger real CUDA init.
    hide_cuda = _is_zerogpu_space()
    old_val = os.environ.get("CUDA_VISIBLE_DEVICES")
    if hide_cuda:
        os.environ["CUDA_VISIBLE_DEVICES"] = ""

    try:
        _converter = DocumentConverter(format_options={InputFormat.PDF: pdf_format_option})
    finally:
        if hide_cuda:
            if old_val is None:
                os.environ.pop("CUDA_VISIBLE_DEVICES", None)
            else:
                os.environ["CUDA_VISIBLE_DEVICES"] = old_val

    return _converter


def parse_document(
    file_bytes: bytes,
    file_ext: str = ".pdf",
    on_progress: Callable[[str], None] | None = None,
) -> dict[str, Any]:
    """Parse a document with Docling and extract markdown, text, and figure regions.

    Args:
        file_bytes: Document file content as bytes.
        file_ext: File extension (e.g. ``".pdf"``, ``".docx"``, ``".xlsx"``, ``".pptx"``).
        on_progress: Optional callback ``(phase_message) -> None`` for progress reporting.

    Returns:
        Dictionary with keys:
            - ``html``: HTML-wrapped markdown representation of the document.
            - ``text``: Full extracted plain text.
            - ``figures``: List of figure dicts with ``bbox``, ``page``, ``caption``, and ``image``.
    """
    def _report(msg: str) -> None:
        if on_progress:
            on_progress(msg)

    try:
        from docling.datamodel.base_models import InputFormat

        ext = file_ext.lower()
        tmp_path = str(get_temp_dir() / f"{uuid.uuid4().hex}{ext}")
        Path(tmp_path).write_bytes(file_bytes)

        try:
            format_name = _EXT_TO_INPUT_FORMAT.get(ext, "PDF")
            input_format = getattr(InputFormat, format_name)

            _report("Initializing document converter...")

            converter = get_converter()

            _report("Converting document (this may take a moment)...")
            result = converter.convert(tmp_path)
            doc = result.document

            _report("Exporting document content...")
            markdown_text = doc.export_to_markdown()
            html = markdown_text
            text = doc.export_to_text()

            _report("Processing figures...")
            figures: list[dict[str, Any]] = []
            try:
                if hasattr(doc, "pictures"):
                    for figure in doc.pictures:
                        if figure.content_layer.value != "body":
                            continue

                        page_num = 0
                        bbox_list = None

                        if figure.prov:
                            page_num = figure.prov[0].page_no - 1  # Docling is 1-based
                            bbox = figure.prov[0].bbox
                            bbox_list = [bbox.l, bbox.t, bbox.width, bbox.height]

                        caption = ""
                        if figure.captions:
                            for cap_ref in figure.captions:
                                try:
                                    if hasattr(cap_ref, "cref") and cap_ref.cref.startswith("#/texts/"):
                                        idx = int(cap_ref.cref.split("/")[-1])
                                        if idx < len(doc.texts):
                                            caption = doc.texts[idx].text
                                            break
                                except Exception:  # noqa: BLE001
                                    pass

                        if figure.image:
                            try:
                                pil_image = figure.image.pil_image
                                figures.append({
                                    "bbox": bbox_list,
                                    "page": page_num,
                                    "caption": caption,
                                    "image": pil_image,
                                })
                            except Exception:  # noqa: BLE001
                                pass

            except Exception:  # noqa: BLE001
                figures = []

            return {"html": html, "text": text, "figures": figures}

        finally:
            if os.path.exists(tmp_path):
                os.unlink(tmp_path)

    except ImportError as e:
        print(f"Docling import error: {e}, using placeholder")
        return {
            "html": "<h1>Sample Document</h1><p>Docling not available - using placeholder.</p>",
            "text": "Sample text from PDF.\n\nDocling not available - using placeholder.",
            "figures": [],
        }
    except Exception as e:  # noqa: BLE001
        import traceback

        print(f"Docling parse error: {e}")
        traceback.print_exc()
        return {
            "html": f"<h1>Error</h1><pre>{e!s}</pre>",
            "text": f"Error: {e!s}",
            "figures": [],
        }