"""
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).
"""

import importlib.util
import shutil
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 .xls_loader import XLSLoader
from .yaml_loader import YAMLLoader


def _module_available(module_name: str) -> bool:
    """Return True if *module_name* can be imported in this environment.

    Uses ``importlib.util.find_spec`` so the (potentially heavy) module is
    located but not actually imported.
    """
    try:
        return importlib.util.find_spec(module_name) is not None
    except (ImportError, ValueError, ModuleNotFoundError):
        # A parent package may itself be missing, which surfaces as
        # ModuleNotFoundError/ImportError from find_spec.
        return False


# Whether a loader *class* can be imported is NOT sufficient to know whether a
# format actually works: the ``unstructured`` loaders import their real parser
# (and its third-party dependency) lazily inside ``.load()``. So a format like
# ODT can be advertised as supported, accept an upload, and only then fail with
# ``ModuleNotFoundError: No module named 'docx'`` (issue #4414).
#
# These flags probe the actual runtime dependency so the registry only
# advertises formats that will really extract text. Missing deps mean the
# format is omitted and the upload path returns a clear "Unsupported format"
# instead of a swallowed extraction failure.
HAS_DOCX_DEP = _module_available("docx")  # python-docx: .docx/.odt
HAS_PPTX_DEP = _module_available("pptx")  # python-pptx: .pptx
# Legacy OLE binary office formats (.doc/.ppt) are NOT read by python-docx/
# python-pptx; unstructured converts them to the modern format by shelling out
# to LibreOffice (``soffice``). Without that binary they fail at runtime with
# "soffice command was not found", so only advertise them when it is present.
HAS_LIBREOFFICE = bool(
    shutil.which("soffice") or shutil.which("libreoffice")
)  # .doc/.ppt conversion
# unstructured.partition.xlsx imports msoffcrypto at module load, so openpyxl
# alone is not enough to extract a modern .xlsx spreadsheet.
HAS_XLSX_DEP = _module_available("openpyxl") and _module_available(
    "msoffcrypto"
)  # openpyxl + msoffcrypto: .xlsx
# Legacy .xls is read by our own XLSLoader (pandas + xlrd), which avoids the
# msoffcrypto pre-check that crashes on some .xls files in the unstructured path.
HAS_XLS_DEP = _module_available("xlrd")  # xlrd: .xls
HAS_PANDOC_DEP = _module_available(
    "pypandoc"
)  # pandoc bridge: epub/rtf/rst/org
HAS_OCR_DEP = _module_available("pytesseract") or _module_available(
    "unstructured.pytesseract"
)  # tesseract OCR for image formats (also needs the tesseract binary)


# 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": {},
    },
    # CSV (Word/Excel are registered below, gated on their runtime deps)
    ".csv": {
        "loader_class": CSVLoader,
        "loader_kwargs": {},
    },
    # HTML
    ".html": {
        "loader_class": UnstructuredHTMLLoader,
        "loader_kwargs": {},
    },
    ".htm": {
        "loader_class": UnstructuredHTMLLoader,
        "loader_kwargs": {},
    },
}

# Word documents - .docx needs python-docx; legacy .doc additionally needs
# LibreOffice to convert the OLE binary to .docx first.
if HAS_DOCX_DEP:
    LOADER_REGISTRY[".docx"] = {
        "loader_class": UnstructuredWordDocumentLoader,
        "loader_kwargs": {},
    }
    if HAS_LIBREOFFICE:
        LOADER_REGISTRY[".doc"] = {
            "loader_class": UnstructuredWordDocumentLoader,
            "loader_kwargs": {},
        }

# Modern spreadsheets - require openpyxl + msoffcrypto (unstructured path)
if HAS_XLSX_DEP:
    LOADER_REGISTRY[".xlsx"] = {
        "loader_class": UnstructuredExcelLoader,
        "loader_kwargs": {},
    }

# Legacy spreadsheets - read directly via our XLSLoader (pandas + xlrd)
if HAS_XLS_DEP:
    LOADER_REGISTRY[".xls"] = {
        "loader_class": XLSLoader,
        "loader_kwargs": {},
    }

# ODT - the unstructured ODT partitioner imports python-docx internally
if HAS_ODT_LOADER and HAS_DOCX_DEP:
    LOADER_REGISTRY[".odt"] = {
        "loader_class": UnstructuredODTLoader,
        "loader_kwargs": {},
    }

# PowerPoint presentations - .pptx needs python-pptx; legacy .ppt additionally
# needs LibreOffice to convert the OLE binary to .pptx first.
if HAS_PPTX_DEP:
    LOADER_REGISTRY[".pptx"] = {
        "loader_class": UnstructuredPowerPointLoader,
        "loader_kwargs": {},
    }
    if HAS_LIBREOFFICE:
        LOADER_REGISTRY[".ppt"] = {
            "loader_class": UnstructuredPowerPointLoader,
            "loader_kwargs": {},
        }

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

# RTF (Rich Text Format) - requires pandoc
if HAS_RTF_LOADER and HAS_PANDOC_DEP:
    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 and HAS_PANDOC_DEP:
    LOADER_REGISTRY[".rst"] = {
        "loader_class": UnstructuredRSTLoader,
        "loader_kwargs": {},
    }

# Org-mode files - requires pandoc
if HAS_ORG_LOADER and HAS_PANDOC_DEP:
    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 - require pytesseract (and the tesseract binary)
if HAS_IMAGE_LOADER and HAS_OCR_DEP:
    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", {})
