"""LangChain document loader for pdfmux.

Usage:
    from langchain_pdfmux import PDFMuxLoader

    loader = PDFMuxLoader("report.pdf")
    docs = loader.load()
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Iterator

from langchain_core.document_loaders import BaseLoader
from langchain_core.documents import Document

logger = logging.getLogger(__name__)

_VALID_QUALITY = {"fast", "standard", "high"}


class PDFMuxLoader(BaseLoader):
    """LangChain document loader powered by pdfmux.

    Extracts text from PDFs using pdfmux's self-healing multi-pass pipeline
    and returns LangChain Document objects with rich metadata.

    Supports single files and directories (all .pdf files in the directory).

    Args:
        file_path: Path to a PDF file or directory containing PDFs.
        quality: Extraction quality — "fast", "standard" (default), or "high".
        glob: Glob pattern for matching PDF files in a directory.
              Defaults to "*.pdf".

    Example::

        from langchain_pdfmux import PDFMuxLoader

        loader = PDFMuxLoader("report.pdf", quality="high")
        docs = loader.load()
        for doc in docs:
            print(doc.metadata["source"], doc.metadata["confidence"])
    """

    def __init__(
        self,
        file_path: str | Path,
        *,
        quality: str = "standard",
        glob: str = "*.pdf",
    ) -> None:
        self.file_path = Path(file_path)
        if quality not in _VALID_QUALITY:
            raise ValueError(
                f"Invalid quality {quality!r}. Must be one of: {', '.join(sorted(_VALID_QUALITY))}"
            )
        self.quality = quality
        self.glob = glob

        # Validate path exists
        if not self.file_path.exists():
            raise ValueError(
                f"Path does not exist: {self.file_path}. "
                "Provide a valid path to a PDF file or directory."
            )

    def lazy_load(self) -> Iterator[Document]:
        """Lazily load documents from the PDF(s).

        Yields:
            Document objects with page_content and metadata.

        Raises:
            ImportError: If pdfmux is not installed.
            RuntimeError: If pdfmux fails to process a file.
        """
        try:
            import pdfmux
        except ImportError:
            raise ImportError(
                "pdfmux is required for PDFMuxLoader. Install it with: pip install pdfmux"
            )

        if self.file_path.is_dir():
            paths = sorted(self.file_path.glob(self.glob))
            if not paths:
                logger.warning("No files matching %r found in %s", self.glob, self.file_path)
        else:
            paths = [self.file_path]

        for path in paths:
            try:
                chunks = pdfmux.load_llm_context(path, quality=self.quality)
            except Exception as exc:
                raise RuntimeError(f"pdfmux failed to process {path}: {exc}") from exc

            for chunk in chunks:
                yield Document(
                    page_content=chunk["text"],
                    metadata={
                        "source": str(path),
                        "title": chunk["title"],
                        "page_start": chunk["page_start"],
                        "page_end": chunk["page_end"],
                        "tokens": chunk["tokens"],
                        "confidence": chunk["confidence"],
                    },
                )
