"""
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] = {}
        # Parallel map of name -> classification metadata, e.g.
        # {"is_local": True}. Kept separate from _retrievers so get()
        # keeps returning the bare retriever object (callers + tests
        # depend on `get() is retriever`).
        self._metadata: Dict[str, dict] = {}
        self._lock = Lock()

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

        Args:
            name: Name to register the retriever under
            retriever: LangChain BaseRetriever instance
            is_local: Whether this retriever reads local/private data
                (a user's vector store / KB) rather than the public
                internet. Defaults to True because programmatically
                registered retrievers are almost always private corpora;
                the egress policy uses this to decide whether the
                retriever may run under PUBLIC_ONLY / PRIVATE_ONLY scopes.
        """
        with self._lock:
            self._retrievers[name] = retriever
            self._metadata[name] = {"is_local": is_local}
            logger.info(
                f"Registered retriever '{name}' of type {type(retriever).__name__} "
                f"(is_local={is_local})"
            )

    def register_multiple(
        self,
        retrievers: Dict[str, "BaseRetriever | dict"],
        is_local: bool = True,
    ) -> None:
        """
        Register multiple retrievers at once.

        Args:
            retrievers: Either ``{name: retriever}`` (uses ``is_local``
                arg as a single default) OR a richer
                ``{name: {"retriever": …, "is_local": bool}}`` mapping
                so different retrievers can carry different
                classifications in one call. Mixing both shapes inside
                the same dict is allowed.
            is_local: Default classification when an entry uses the
                bare-retriever shape. The public programmatic API
                (api/research_functions.py) calls this; we default to
                True here ONLY because the old API contract was True.
                CLI / programmatic callers passing a cloud-backed
                retriever MUST use the dict shape with
                ``is_local=False`` or they will egress under
                PRIVATE_ONLY (silent classification lie).
        """
        with self._lock:
            for name, entry in retrievers.items():
                if isinstance(entry, dict):
                    retriever = entry.get("retriever")
                    entry_is_local = bool(entry.get("is_local", is_local))
                else:
                    retriever = entry
                    entry_is_local = is_local
                if retriever is None:
                    logger.warning(
                        f"Retriever entry '{name}' has no retriever — skipped"
                    )
                    continue
                self._retrievers[name] = retriever
                self._metadata[name] = {"is_local": entry_is_local}
                logger.info(
                    f"Registered retriever '{name}' of type {type(retriever).__name__} "
                    f"(is_local={entry_is_local})"
                )

    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]
                self._metadata.pop(name, None)
                logger.info(f"Unregistered retriever '{name}'")

    def clear(self) -> None:
        """Clear all registered retrievers."""
        with self._lock:
            count = len(self._retrievers)
            self._retrievers.clear()
            self._metadata.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())

    def get_metadata(
        self, name: str, username: Optional[str] = None
    ) -> Optional[Dict]:
        """Return policy-relevant metadata for a registered retriever.

        Returns a dict like ``{"is_local": True}`` for a registered
        retriever, or ``None`` when the retriever is unknown. A
        registered retriever with no recorded classification yields an
        empty dict, which ``evaluate_retriever`` treats as "unclassified"
        and fails closed under any non-BOTH scope.

        The egress policy's ``evaluate_retriever`` consults this hook to
        decide whether a retriever may run under the active scope.

        ``username`` is accepted so future per-user retriever isolation
        can plug in without changing the call sites.
        """
        with self._lock:
            if name not in self._retrievers:
                return None
            # Copy so callers can't mutate the stored classification.
            return dict(self._metadata.get(name, {}))


# Global registry instance
retriever_registry = RetrieverRegistry()
