"""
Registry for dynamically registering LangChain retrievers as search engines.
"""

from threading import Lock
from typing import Dict, Optional

from langchain_core.retrievers import BaseRetriever
from loguru import logger


class RetrieverRegistry:
    """
    Thread-safe registry for LangChain retrievers.

    This allows users to register retrievers programmatically and use them
    as search engines within LDR.
    """

    def __init__(self):
        self._retrievers: Dict[str, BaseRetriever] = {}
        self._lock = Lock()

    def register(self, name: str, retriever: BaseRetriever) -> None:
        """
        Register a retriever with a given name.

        Args:
            name: Name to register the retriever under
            retriever: LangChain BaseRetriever instance
        """
        with self._lock:
            self._retrievers[name] = retriever
            logger.info(
                f"Registered retriever '{name}' of type {type(retriever).__name__}"
            )

    def register_multiple(self, retrievers: Dict[str, BaseRetriever]) -> None:
        """
        Register multiple retrievers at once.

        Args:
            retrievers: Dictionary of {name: retriever} pairs
        """
        with self._lock:
            for name, retriever in retrievers.items():
                self._retrievers[name] = retriever
                logger.info(
                    f"Registered retriever '{name}' of type {type(retriever).__name__}"
                )

    def get(self, name: str) -> Optional[BaseRetriever]:
        """
        Get a registered retriever by name.

        Args:
            name: Name of the retriever

        Returns:
            The retriever if found, None otherwise
        """
        with self._lock:
            return self._retrievers.get(name)

    def unregister(self, name: str) -> None:
        """
        Remove a registered retriever.

        Args:
            name: Name of the retriever to remove
        """
        with self._lock:
            if name in self._retrievers:
                del self._retrievers[name]
                logger.info(f"Unregistered retriever '{name}'")

    def clear(self) -> None:
        """Clear all registered retrievers."""
        with self._lock:
            count = len(self._retrievers)
            self._retrievers.clear()
            logger.info(f"Cleared {count} registered retrievers")

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

        Args:
            name: Name of the retriever

        Returns:
            True if registered, False otherwise
        """
        with self._lock:
            return name in self._retrievers

    def list_registered(self) -> list[str]:
        """
        Get list of all registered retriever names.

        Returns:
            List of retriever names
        """
        with self._lock:
            return list(self._retrievers.keys())


# Global registry instance
retriever_registry = RetrieverRegistry()
