"""Tests for the retriever registry."""

from langchain.schema import BaseRetriever, Document

from local_deep_research.web_search_engines.retriever_registry import (
    RetrieverRegistry,
)


class MockRetriever(BaseRetriever):
    """Mock retriever for testing."""

    name: str = "mock"

    class Config:
        """Pydantic config."""

        arbitrary_types_allowed = True

    def _get_relevant_documents(self, query: str):
        return [Document(page_content=f"Mock result from {self.name}")]

    async def _aget_relevant_documents(self, query: str):
        return self._get_relevant_documents(query)


class TestRetrieverRegistry:
    """Test the retriever registry functionality."""

    def test_register_single_retriever(self):
        """Test registering a single retriever."""
        registry = RetrieverRegistry()
        retriever = MockRetriever(name="test1")

        registry.register("test1", retriever)

        assert registry.is_registered("test1")
        assert registry.get("test1") == retriever
        assert "test1" in registry.list_registered()

    def test_register_multiple_retrievers(self):
        """Test registering multiple retrievers at once."""
        registry = RetrieverRegistry()
        retrievers = {
            "test1": MockRetriever(name="test1"),
            "test2": MockRetriever(name="test2"),
            "test3": MockRetriever(name="test3"),
        }

        registry.register_multiple(retrievers)

        assert len(registry.list_registered()) == 3
        assert all(name in registry.list_registered() for name in retrievers)
        assert registry.get("test2") == retrievers["test2"]

    def test_unregister_retriever(self):
        """Test unregistering a retriever."""
        registry = RetrieverRegistry()
        retriever = MockRetriever()

        registry.register("test", retriever)
        assert registry.is_registered("test")

        registry.unregister("test")
        assert not registry.is_registered("test")
        assert registry.get("test") is None

    def test_clear_registry(self):
        """Test clearing all retrievers."""
        registry = RetrieverRegistry()

        # Register multiple retrievers
        for i in range(5):
            registry.register(f"test{i}", MockRetriever(name=f"test{i}"))

        assert len(registry.list_registered()) == 5

        registry.clear()
        assert len(registry.list_registered()) == 0

    def test_get_nonexistent_retriever(self):
        """Test getting a retriever that doesn't exist."""
        registry = RetrieverRegistry()
        assert registry.get("nonexistent") is None
        assert not registry.is_registered("nonexistent")

    def test_thread_safety(self):
        """Test thread safety of registry operations."""
        import threading
        import time

        registry = RetrieverRegistry()

        def register_retrievers(start_idx):
            for i in range(10):
                name = f"retriever_{start_idx}_{i}"
                registry.register(name, MockRetriever(name=name))
                time.sleep(0.001)  # Small delay to increase chance of conflicts

        # Create multiple threads
        threads = []
        for i in range(5):
            thread = threading.Thread(
                target=register_retrievers, args=(i * 10,)
            )
            threads.append(thread)
            thread.start()

        # Wait for all threads to complete
        for thread in threads:
            thread.join()

        # Verify all retrievers were registered
        assert len(registry.list_registered()) == 50

    def test_overwrite_existing_retriever(self):
        """Test that registering with same name overwrites."""
        registry = RetrieverRegistry()

        retriever1 = MockRetriever(name="first")
        retriever2 = MockRetriever(name="second")

        registry.register("test", retriever1)
        assert registry.get("test").name == "first"

        registry.register("test", retriever2)
        assert registry.get("test").name == "second"
