"""
Document loader registry - maps file extensions to LangChain loaders.

This module provides a centralized registry for document loaders that can be used
by both collection uploads (bytes) and local search (file paths).
"""

from pathlib import Path
from typing import Optional

from langchain_community.document_loaders import (
    CSVLoader,
    EverNoteLoader,
    MHTMLLoader,
    NotebookLoader,
    PyPDFLoader,
    TextLoader,
    TomlLoader,
    UnstructuredEmailLoader,
    UnstructuredExcelLoader,
    UnstructuredHTMLLoader,
    UnstructuredMarkdownLoader,
    UnstructuredPowerPointLoader,
    UnstructuredWordDocumentLoader,
    UnstructuredXMLLoader,
)
from langchain_core.document_loaders import BaseLoader
from loguru import logger

# Import loaders that require external system tools (pandoc, tesseract, etc.)
# These may fail at runtime if the tools aren't installed
try:
    from langchain_community.document_loaders import UnstructuredODTLoader

    HAS_ODT_LOADER = True
except ImportError:
    HAS_ODT_LOADER = False
    logger.debug("UnstructuredODTLoader not available - ODT support disabled")

try:
    from langchain_community.document_loaders import UnstructuredEPubLoader

    HAS_EPUB_LOADER = True
except ImportError:
    HAS_EPUB_LOADER = False
    logger.debug("UnstructuredEPubLoader not available - EPUB support disabled")

try:
    from langchain_community.document_loaders import UnstructuredRTFLoader

    HAS_RTF_LOADER = True
except ImportError:
    HAS_RTF_LOADER = False
    logger.debug("UnstructuredRTFLoader not available - RTF support disabled")

try:
    from langchain_community.document_loaders import UnstructuredRSTLoader

    HAS_RST_LOADER = True
except ImportError:
    HAS_RST_LOADER = False
    logger.debug("UnstructuredRSTLoader not available - RST support disabled")

try:
    from langchain_community.document_loaders import UnstructuredOrgModeLoader

    HAS_ORG_LOADER = True
except ImportError:
    HAS_ORG_LOADER = False
    logger.debug(
        "UnstructuredOrgModeLoader not available - Org support disabled"
    )

try:
    from langchain_community.document_loaders import UnstructuredImageLoader

    HAS_IMAGE_LOADER = True
except ImportError:
    HAS_IMAGE_LOADER = False
    logger.debug(
        "UnstructuredImageLoader not available - Image/OCR support disabled"
    )

# Import our custom loaders
from .json_loader import SimpleJSONLoader
from .yaml_loader import YAMLLoader


# Extension to loader mapping
# Each entry contains:
#   - loader_class: The LangChain loader class
#   - loader_kwargs: Optional kwargs to pass to the loader
#   - requires_path: Whether the loader requires a file path (vs bytes)
LOADER_REGISTRY: dict = {
    # PDF
    ".pdf": {
        "loader_class": PyPDFLoader,
        "loader_kwargs": {},
    },
    # Plain text
    ".txt": {
        "loader_class": TextLoader,
        "loader_kwargs": {"encoding": "utf-8", "autodetect_encoding": True},
    },
    # Markdown
    ".md": {
        "loader_class": UnstructuredMarkdownLoader,
        "loader_kwargs": {},
    },
    ".markdown": {
        "loader_class": UnstructuredMarkdownLoader,
        "loader_kwargs": {},
    },
    # Word documents
    ".docx": {
        "loader_class": UnstructuredWordDocumentLoader,
        "loader_kwargs": {},
    },
    ".doc": {
        "loader_class": UnstructuredWordDocumentLoader,
        "loader_kwargs": {},
    },
    # Spreadsheets
    ".csv": {
        "loader_class": CSVLoader,
        "loader_kwargs": {},
    },
    ".xlsx": {
        "loader_class": UnstructuredExcelLoader,
        "loader_kwargs": {},
    },
    ".xls": {
        "loader_class": UnstructuredExcelLoader,
        "loader_kwargs": {},
    },
    # HTML
    ".html": {
        "loader_class": UnstructuredHTMLLoader,
        "loader_kwargs": {},
    },
    ".htm": {
        "loader_class": UnstructuredHTMLLoader,
        "loader_kwargs": {},
    },
}

# ODT (requires pypandoc)
if HAS_ODT_LOADER:
    LOADER_REGISTRY[".odt"] = {
        "loader_class": UnstructuredODTLoader,
        "loader_kwargs": {},
    }

# PowerPoint presentations
LOADER_REGISTRY[".ppt"] = {
    "loader_class": UnstructuredPowerPointLoader,
    "loader_kwargs": {},
}
LOADER_REGISTRY[".pptx"] = {
    "loader_class": UnstructuredPowerPointLoader,
    "loader_kwargs": {},
}

# EPUB (ebooks, technical manuals) - requires pandoc
if HAS_EPUB_LOADER:
    LOADER_REGISTRY[".epub"] = {
        "loader_class": UnstructuredEPubLoader,
        "loader_kwargs": {},
    }

# RTF (Rich Text Format) - requires pandoc
if HAS_RTF_LOADER:
    LOADER_REGISTRY[".rtf"] = {
        "loader_class": UnstructuredRTFLoader,
        "loader_kwargs": {},
    }

# XML (important for USPTO patent data)
LOADER_REGISTRY[".xml"] = {
    "loader_class": UnstructuredXMLLoader,
    "loader_kwargs": {},
}

# RST (reStructuredText) - requires pandoc
if HAS_RST_LOADER:
    LOADER_REGISTRY[".rst"] = {
        "loader_class": UnstructuredRSTLoader,
        "loader_kwargs": {},
    }

# Org-mode files - requires pandoc
if HAS_ORG_LOADER:
    LOADER_REGISTRY[".org"] = {
        "loader_class": UnstructuredOrgModeLoader,
        "loader_kwargs": {},
    }

# Email files
LOADER_REGISTRY[".eml"] = {
    "loader_class": UnstructuredEmailLoader,
    "loader_kwargs": {},
}

# TSV (Tab-Separated Values) - use CSVLoader with tab delimiter
LOADER_REGISTRY[".tsv"] = {
    "loader_class": CSVLoader,
    "loader_kwargs": {"csv_args": {"delimiter": "\t"}},
}

# JSON loader using our custom SimpleJSONLoader (no jq dependency)
LOADER_REGISTRY[".json"] = {
    "loader_class": SimpleJSONLoader,
    "loader_kwargs": {},
}

# YAML loader using our custom YAMLLoader
LOADER_REGISTRY[".yaml"] = {
    "loader_class": YAMLLoader,
    "loader_kwargs": {},
}
LOADER_REGISTRY[".yml"] = {
    "loader_class": YAMLLoader,
    "loader_kwargs": {},
}

# Images with OCR support (requires tesseract)
if HAS_IMAGE_LOADER:
    for ext in [".png", ".jpg", ".jpeg", ".tiff", ".tif", ".bmp", ".heic"]:
        LOADER_REGISTRY[ext] = {
            "loader_class": UnstructuredImageLoader,
            "loader_kwargs": {},
        }

# Jupyter notebooks (research code, data analysis)
LOADER_REGISTRY[".ipynb"] = {
    "loader_class": NotebookLoader,
    "loader_kwargs": {
        "include_outputs": True,
        "remove_newline": True,
    },
}

# Evernote exports
LOADER_REGISTRY[".enex"] = {
    "loader_class": EverNoteLoader,
    "loader_kwargs": {"load_single_document": False},
}

# TOML config files
LOADER_REGISTRY[".toml"] = {
    "loader_class": TomlLoader,
    "loader_kwargs": {},
}

# MHTML web archives (saved web pages)
LOADER_REGISTRY[".mhtml"] = {
    "loader_class": MHTMLLoader,
    "loader_kwargs": {},
}
LOADER_REGISTRY[".mht"] = {
    "loader_class": MHTMLLoader,
    "loader_kwargs": {},
}


def get_supported_extensions() -> list[str]:
    """Get list of all supported file extensions."""
    return list(LOADER_REGISTRY.keys())


def is_extension_supported(extension: str) -> bool:
    """Check if a file extension is supported."""
    ext = (
        extension.lower()
        if extension.startswith(".")
        else f".{extension.lower()}"
    )
    return ext in LOADER_REGISTRY


def get_loader_for_path(file_path: str | Path) -> Optional[BaseLoader]:
    """
    Get an appropriate document loader for a file based on its extension.

    Args:
        file_path: Path to the file to load

    Returns:
        A LangChain BaseLoader instance, or None if the extension is not supported
    """
    file_path = Path(file_path)
    extension = file_path.suffix.lower()

    loader_info = get_loader_class_for_extension(extension)
    if loader_info is None:
        logger.warning(
            f"Unsupported file extension: {extension} for {file_path}"
        )
        return None

    loader_class, loader_kwargs = loader_info

    try:
        return loader_class(str(file_path), **loader_kwargs)  # type: ignore[no-any-return]
    except Exception:
        logger.exception(f"Error creating loader for {file_path}")
        return None


def get_loader_class_for_extension(
    extension: str,
) -> Optional[tuple[type, dict]]:
    """
    Get the loader class and kwargs for an extension.

    Args:
        extension: File extension (with or without leading dot)

    Returns:
        Tuple of (loader_class, loader_kwargs) or None if not supported
    """
    ext = (
        extension.lower()
        if extension.startswith(".")
        else f".{extension.lower()}"
    )

    if ext not in LOADER_REGISTRY:
        return None

    entry = LOADER_REGISTRY[ext]
    return entry["loader_class"], entry.get("loader_kwargs", {})
