src/ingestion/docling_loader.py

"""
Docling-based document loader for multi-format document processing.

Supports: PDF, DOCX, PPTX, HTML, images, and Markdown.
Provides structure-aware parsing with table extraction and hierarchy preservation.
"""

import os
import glob
from pathlib import Path
from typing import List, Dict, Optional, Any
from dataclasses import dataclass, field
import logging

logger = logging.getLogger(__name__)

# Supported file extensions
SUPPORTED_EXTENSIONS = {
    ".pdf", ".docx", ".pptx", ".xlsx",
    ".html", ".htm",
    ".md", ".markdown",
    ".png", ".jpg", ".jpeg", ".tiff", ".bmp"
}


@dataclass
class DocumentElement:
    """Represents a structural element in a document."""
    element_type: str  # paragraph, table, heading, list, code, image
    text: str
    level: int = 0  # heading level (1-6) or nesting depth
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class ParsedDocument:
    """Result of parsing a document with Docling."""
    filename: str
    path: str
    elements: List[DocumentElement]
    format: str
    page_count: int = 0
    metadata: Dict[str, Any] = field(default_factory=dict)
    status: str = "OK"
    error: Optional[str] = None

    @property
    def full_text(self) -> str:
        """Get concatenated text from all elements."""
        return "\n\n".join(el.text for el in self.elements if el.text.strip())

    @property
    def chars(self) -> int:
        return len(self.full_text)

    @property
    def words(self) -> int:
        return len(self.full_text.split())


def _get_docling_converter():
    """Lazy load Docling converter to avoid import overhead."""
    try:
        from docling.document_converter import DocumentConverter
        return DocumentConverter()
    except ImportError as e:
        logger.error(f"Docling not installed: {e}")
        raise ImportError(
            "Docling is required for multi-format document loading. "
            "Install with: pip install docling"
        ) from e


def _extract_elements_from_docling(doc_result) -> List[DocumentElement]:
    """
    Extract structured elements from a Docling conversion result.

    Args:
        doc_result: Docling ConversionResult object

    Returns:
        List of DocumentElement objects
    """
    elements = []

    try:
        # Get the DoclingDocument
        docling_doc = doc_result.document

        # Iterate through document items
        for item, level in docling_doc.iterate_items():
            item_type = item.__class__.__name__.lower()

            # Map Docling item types to our element types
            if "heading" in item_type or "title" in item_type:
                el_type = "heading"
                el_level = getattr(item, "level", 1)
            elif "table" in item_type:
                el_type = "table"
                el_level = 0
            elif "list" in item_type:
                el_type = "list"
                el_level = level
            elif "code" in item_type:
                el_type = "code"
                el_level = 0
            elif "image" in item_type or "figure" in item_type:
                el_type = "image"
                el_level = 0
            else:
                el_type = "paragraph"
                el_level = level

            # Extract text content
            text = ""
            if hasattr(item, "text") and item.text:
                text = item.text
            elif hasattr(item, "export_to_markdown"):
                try:
                    # Some items require doc parameter
                    text = item.export_to_markdown(docling_doc)
                except TypeError:
                    try:
                        text = item.export_to_markdown()
                    except Exception:
                        text = str(item) if hasattr(item, "__str__") else ""
            elif hasattr(item, "__str__"):
                text = str(item)

            if text and text.strip():
                elements.append(DocumentElement(
                    element_type=el_type,
                    text=text.strip(),
                    level=el_level,
                    metadata={
                        "original_type": item_type,
                        "depth": level
                    }
                ))

    except Exception as e:
        logger.warning(f"Error extracting elements: {e}")
        # Fallback: try to get markdown export
        try:
            md_text = doc_result.document.export_to_markdown()
            if md_text:
                elements.append(DocumentElement(
                    element_type="paragraph",
                    text=md_text,
                    level=0
                ))
        except Exception:
            pass

    return elements


def load_document_with_docling(file_path: str) -> ParsedDocument:
    """
    Load a single document using Docling.

    Args:
        file_path: Path to the document file

    Returns:
        ParsedDocument with extracted structure and content
    """
    path = Path(file_path)

    if not path.exists():
        return ParsedDocument(
            filename=path.name,
            path=str(path),
            elements=[],
            format=path.suffix.lower(),
            status="ERROR",
            error=f"File not found: {file_path}"
        )

    ext = path.suffix.lower()
    if ext not in SUPPORTED_EXTENSIONS:
        return ParsedDocument(
            filename=path.name,
            path=str(path),
            elements=[],
            format=ext,
            status="SKIPPED",
            error=f"Unsupported format: {ext}"
        )

    try:
        converter = _get_docling_converter()
        result = converter.convert(str(path))

        elements = _extract_elements_from_docling(result)

        # Get page count if available
        page_count = 0
        try:
            if hasattr(result.document, "pages"):
                page_count = len(result.document.pages)
        except Exception:
            pass

        return ParsedDocument(
            filename=path.name,
            path=str(path),
            elements=elements,
            format=ext,
            page_count=page_count,
            metadata={
                "converter": "docling",
                "element_count": len(elements)
            },
            status="OK"
        )

    except Exception as e:
        logger.error(f"Error processing {file_path}: {e}")
        return ParsedDocument(
            filename=path.name,
            path=str(path),
            elements=[],
            format=ext,
            status="ERROR",
            error=str(e)
        )


def load_documents_with_docling(
    dir_path: str,
    extensions: Optional[List[str]] = None,
    max_chars: int = 50000,
    recursive: bool = False
) -> List[ParsedDocument]:
    """
    Load multiple documents from a directory using Docling.

    Args:
        dir_path: Path to directory containing documents
        extensions: List of extensions to process (default: all supported)
        max_chars: Maximum characters per document (skip larger files)
        recursive: Whether to search subdirectories

    Returns:
        List of ParsedDocument objects
    """
    path = Path(dir_path).expanduser()

    if not path.is_dir():
        raise FileNotFoundError(f"Directory not found: {dir_path}")

    if extensions is None:
        extensions = list(SUPPORTED_EXTENSIONS)
    else:
        extensions = [e if e.startswith(".") else f".{e}" for e in extensions]

    # Find all matching files
    files = []
    for ext in extensions:
        pattern = f"**/*{ext}" if recursive else f"*{ext}"
        files.extend(path.glob(pattern))

    files = sorted(set(files))

    documents = []
    for file_path in files:
        doc = load_document_with_docling(str(file_path))

        # Check size limit
        if doc.status == "OK" and doc.chars > max_chars:
            doc.status = "SKIPPED_TOO_LARGE"
            doc.error = f"Document exceeds {max_chars} chars ({doc.chars})"
            doc.elements = []

        documents.append(doc)

    return documents


def convert_to_legacy_format(docs: List[ParsedDocument]) -> List[Dict]:
    """
    Convert ParsedDocument list to legacy format for backward compatibility.

    Args:
        docs: List of ParsedDocument objects

    Returns:
        List of dicts matching load_markdown_docs output format
    """
    legacy = []
    for doc in docs:
        legacy.append({
            "filename": doc.filename,
            "path": doc.path,
            "text": doc.full_text if doc.status == "OK" else None,
            "chars": doc.chars,
            "words": doc.words,
            "status": doc.status,
            "format": doc.format,
            "elements": doc.elements,  # Additional: structured elements
            "page_count": doc.page_count,
            "metadata": doc.metadata
        })
    return legacy


def print_summary(docs: List[ParsedDocument]):
    """Print summary of loaded documents."""
    if not docs:
        print("No documents found or all were skipped.")
        return

    print(f"{'FILENAME':40} {'FORMAT':8} {'STATUS':20} {'CHARS':>8} {'ELEMENTS':>8}")
    print("-" * 90)

    for d in docs:
        name = d.filename[:40]
        fmt = d.format[:8]
        status = d.status[:20]
        chars = d.chars
        elements = len(d.elements)
        print(f"{name:40} {fmt:8} {status:20} {chars:8d} {elements:8d}")

    ok_count = sum(1 for d in docs if d.status == "OK")
    skipped = len(docs) - ok_count
    print("-" * 90)
    print(f"Total: {len(docs)}  OK: {ok_count}  Skipped/Errors: {skipped}")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Load documents using Docling for RAG ingestion."
    )
    parser.add_argument("dir", help="Directory containing documents")
    parser.add_argument(
        "--extensions", "-e",
        nargs="+",
        default=None,
        help="File extensions to process (default: all supported)"
    )
    parser.add_argument(
        "--max-chars",
        type=int,
        default=50000,
        help="Max characters to accept (default: 50000)"
    )
    parser.add_argument(
        "--recursive", "-r",
        action="store_true",
        help="Search subdirectories recursively"
    )

    args = parser.parse_args()

    docs = load_documents_with_docling(
        args.dir,
        extensions=args.extensions,
        max_chars=args.max_chars,
        recursive=args.recursive
    )
    print_summary(docs)