"""Registry for custom LangChain LLMs.

This module provides a global registry for registering and managing custom LangChain
LLMs that can be used with Local Deep Research.
"""

import threading
from typing import Callable, Dict, Optional, Union

from langchain.chat_models.base import BaseChatModel
from loguru import logger


class LLMRegistry:
    """Thread-safe registry for custom LangChain LLMs."""

    def __init__(self):
        self._llms: Dict[
            str, Union[BaseChatModel, Callable[..., BaseChatModel]]
        ] = {}
        self._lock = threading.Lock()

    def register(
        self, name: str, llm: Union[BaseChatModel, Callable[..., BaseChatModel]]
    ) -> None:
        """Register a custom LLM.

        Args:
            name: Unique name for the LLM (case-insensitive)
            llm: Either a BaseChatModel instance or a factory function that returns one
        """
        with self._lock:
            # Normalize name to lowercase for case-insensitive storage
            normalized_name = name.lower()
            if normalized_name in self._llms:
                logger.warning(f"Overwriting existing LLM: {name}")
            self._llms[normalized_name] = llm
            logger.info(
                f"Registered custom LLM: {name} (normalized: {normalized_name})"
            )

    # Completes the CRUD API surface for the registry.
    # Used in tests to verify cleanup behavior.
    def unregister(self, name: str) -> None:
        """Unregister a custom LLM.

        Args:
            name: Name of the LLM to unregister (case-insensitive)
        """
        with self._lock:
            normalized_name = name.lower()
            if normalized_name in self._llms:
                del self._llms[normalized_name]
                logger.info(f"Unregistered custom LLM: {name}")

    def get(
        self, name: str
    ) -> Optional[Union[BaseChatModel, Callable[..., BaseChatModel]]]:
        """Get a registered LLM.

        Args:
            name: Name of the LLM to retrieve (case-insensitive)

        Returns:
            The LLM instance/factory or None if not found
        """
        with self._lock:
            normalized_name = name.lower()
            return self._llms.get(normalized_name)

    def is_registered(self, name: str) -> bool:
        """Check if an LLM is registered.

        Args:
            name: Name to check (case-insensitive)

        Returns:
            True if registered, False otherwise
        """
        with self._lock:
            normalized_name = name.lower()
            return normalized_name in self._llms

    # Used in test assertions to verify registry state;
    # part of public API for plugin authors.
    def list_registered(self) -> list[str]:
        """Get list of all registered LLM names.

        Returns:
            List of registered LLM names
        """
        with self._lock:
            return list(self._llms.keys())

    # Used in 7+ test files' autouse fixtures for test isolation
    # (64+ tests depend on this to reset global state between runs).
    def clear(self) -> None:
        """Clear all registered LLMs."""
        with self._lock:
            self._llms.clear()
            logger.info("Cleared all registered custom LLMs")


# Global registry instance
_llm_registry = LLMRegistry()


# Public API functions
def register_llm(
    name: str, llm: Union[BaseChatModel, Callable[..., BaseChatModel]]
) -> None:
    """Register a custom LLM in the global registry.

    Args:
        name: Unique name for the LLM
        llm: Either a BaseChatModel instance or a factory function
    """
    _llm_registry.register(name, llm)


def unregister_llm(name: str) -> None:
    """Unregister a custom LLM from the global registry.

    Args:
        name: Name of the LLM to unregister
    """
    _llm_registry.unregister(name)


def get_llm_from_registry(
    name: str,
) -> Optional[Union[BaseChatModel, Callable[..., BaseChatModel]]]:
    """Get a registered LLM from the global registry.

    Args:
        name: Name of the LLM to retrieve

    Returns:
        The LLM instance/factory or None if not found
    """
    return _llm_registry.get(name)


def is_llm_registered(name: str) -> bool:
    """Check if an LLM is registered in the global registry.

    Args:
        name: Name to check

    Returns:
        True if registered, False otherwise
    """
    return _llm_registry.is_registered(name)


def list_registered_llms() -> list[str]:
    """Get list of all registered LLM names.

    Returns:
        List of registered LLM names
    """
    return _llm_registry.list_registered()


def clear_llm_registry() -> None:
    """Clear all registered LLMs from the global registry."""
    _llm_registry.clear()
