"""Tests for retrieve_user_context tool."""

import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from redis.exceptions import RedisError
from pymongo.errors import PyMongoError

from src.geaicontext.tools.retrieve import retrieve_user_context


class TestRetrieveUserContext:
    """Test cases for retrieve_user_context function."""

    @pytest.fixture
    def mock_redis_client_instance(self):
        """Create a mock RedisClient instance."""
        client = MagicMock()
        client.connect = AsyncMock()
        client.disconnect = AsyncMock()
        client.get_context = AsyncMock(return_value=None)
        client.set_context = AsyncMock(return_value=True)
        return client

    @pytest.fixture
    def mock_mongo_client_instance(self):
        """Create a mock MongoClient instance."""
        client = MagicMock()
        client.connect = AsyncMock()
        client.disconnect = AsyncMock()
        client.get_context = AsyncMock(return_value=None)
        return client

    @pytest.mark.asyncio
    async def test_retrieve_empty_customer_id(self):
        """Test retrieve with empty customer_id."""
        result = await retrieve_user_context("")
        assert result is None

    @pytest.mark.asyncio
    async def test_retrieve_cache_hit(
        self, mock_redis_client_instance, mock_mongo_client_instance, sample_context, sample_customer_id
    ):
        """Test retrieve when context is found in cache."""
        mock_redis_client_instance.get_context = AsyncMock(return_value=sample_context)
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis_client_instance)
        ), patch(
            "src.geaicontext.tools.retrieve.get_mongo_client",
            AsyncMock(return_value=mock_mongo_client_instance)
        ):
            result = await retrieve_user_context(sample_customer_id)
            
            assert result is not None
            assert result["name"] == "John Doe"
            mock_redis_client_instance.get_context.assert_called_once_with(sample_customer_id)
            # MongoDB should not be called on cache hit
            mock_mongo_client_instance.get_context.assert_not_called()

    @pytest.mark.asyncio
    async def test_retrieve_cache_miss_db_hit(
        self, mock_redis_client_instance, mock_mongo_client_instance, sample_context, sample_customer_id
    ):
        """Test retrieve when cache miss but found in database."""
        mock_redis_client_instance.get_context = AsyncMock(return_value=None)
        mock_mongo_client_instance.get_context = AsyncMock(return_value=sample_context)
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis_client_instance)
        ), patch(
            "src.geaicontext.tools.retrieve.get_mongo_client",
            AsyncMock(return_value=mock_mongo_client_instance)
        ):
            result = await retrieve_user_context(sample_customer_id)
            
            assert result is not None
            assert result["name"] == "John Doe"
            # Should populate cache after DB hit
            mock_redis_client_instance.set_context.assert_called_once_with(sample_customer_id, sample_context)

    @pytest.mark.asyncio
    async def test_retrieve_cache_miss_db_miss(
        self, mock_redis_client_instance, mock_mongo_client_instance, sample_customer_id
    ):
        """Test retrieve when not found in cache or database."""
        mock_redis_client_instance.get_context = AsyncMock(return_value=None)
        mock_mongo_client_instance.get_context = AsyncMock(return_value=None)
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis_client_instance)
        ), patch(
            "src.geaicontext.tools.retrieve.get_mongo_client",
            AsyncMock(return_value=mock_mongo_client_instance)
        ):
            result = await retrieve_user_context(sample_customer_id)
            
            assert result is None

    @pytest.mark.asyncio
    async def test_retrieve_redis_error_fallback_to_db(
        self, mock_redis_client_instance, mock_mongo_client_instance, sample_context, sample_customer_id
    ):
        """Test retrieve falls back to DB when Redis fails."""
        mock_redis_client_instance.get_context = AsyncMock(side_effect=RedisError("Redis down"))
        mock_mongo_client_instance.get_context = AsyncMock(return_value=sample_context)
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis_client_instance)
        ), patch(
            "src.geaicontext.tools.retrieve.get_mongo_client",
            AsyncMock(return_value=mock_mongo_client_instance)
        ):
            result = await retrieve_user_context(sample_customer_id)
            
            assert result is not None
            assert result["name"] == "John Doe"

    @pytest.mark.asyncio
    async def test_retrieve_redis_unexpected_error_fallback_to_db(
        self, mock_redis_client_instance, mock_mongo_client_instance, sample_context, sample_customer_id
    ):
        """Test retrieve falls back to DB when Redis has unexpected error."""
        mock_redis_client_instance.get_context = AsyncMock(side_effect=Exception("Unexpected"))
        mock_mongo_client_instance.get_context = AsyncMock(return_value=sample_context)
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis_client_instance)
        ), patch(
            "src.geaicontext.tools.retrieve.get_mongo_client",
            AsyncMock(return_value=mock_mongo_client_instance)
        ):
            result = await retrieve_user_context(sample_customer_id)
            
            assert result is not None

    @pytest.mark.asyncio
    async def test_retrieve_mongodb_error(
        self, mock_redis_client_instance, mock_mongo_client_instance, sample_customer_id
    ):
        """Test retrieve returns None when MongoDB fails."""
        mock_redis_client_instance.get_context = AsyncMock(return_value=None)
        mock_mongo_client_instance.get_context = AsyncMock(side_effect=PyMongoError("MongoDB down"))
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis_client_instance)
        ), patch(
            "src.geaicontext.tools.retrieve.get_mongo_client",
            AsyncMock(return_value=mock_mongo_client_instance)
        ):
            result = await retrieve_user_context(sample_customer_id)
            
            assert result is None

    @pytest.mark.asyncio
    async def test_retrieve_mongodb_unexpected_error(
        self, mock_redis_client_instance, mock_mongo_client_instance, sample_customer_id
    ):
        """Test retrieve returns None when MongoDB has unexpected error."""
        mock_redis_client_instance.get_context = AsyncMock(return_value=None)
        mock_mongo_client_instance.get_context = AsyncMock(side_effect=Exception("Unexpected"))
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis_client_instance)
        ), patch(
            "src.geaicontext.tools.retrieve.get_mongo_client",
            AsyncMock(return_value=mock_mongo_client_instance)
        ):
            result = await retrieve_user_context(sample_customer_id)
            
            assert result is None

    @pytest.mark.asyncio
    async def test_retrieve_cache_populate_fails(
        self, mock_redis_client_instance, mock_mongo_client_instance, sample_context, sample_customer_id
    ):
        """Test retrieve succeeds even when cache population fails."""
        mock_redis_client_instance.get_context = AsyncMock(return_value=None)
        mock_redis_client_instance.set_context = AsyncMock(side_effect=RedisError("Cache write failed"))
        mock_mongo_client_instance.get_context = AsyncMock(return_value=sample_context)
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis_client_instance)
        ), patch(
            "src.geaicontext.tools.retrieve.get_mongo_client",
            AsyncMock(return_value=mock_mongo_client_instance)
        ):
            result = await retrieve_user_context(sample_customer_id)
            
            # Should still return the context even if cache population fails
            assert result is not None
            assert result["name"] == "John Doe"

    @pytest.mark.asyncio
    async def test_retrieve_cache_populate_unexpected_error(
        self, mock_redis_client_instance, mock_mongo_client_instance, sample_context, sample_customer_id
    ):
        """Test retrieve succeeds even when cache population has unexpected error."""
        mock_redis_client_instance.get_context = AsyncMock(return_value=None)
        mock_redis_client_instance.set_context = AsyncMock(side_effect=Exception("Unexpected"))
        mock_mongo_client_instance.get_context = AsyncMock(return_value=sample_context)
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis_client_instance)
        ), patch(
            "src.geaicontext.tools.retrieve.get_mongo_client",
            AsyncMock(return_value=mock_mongo_client_instance)
        ):
            result = await retrieve_user_context(sample_customer_id)
            
            assert result is not None

    @pytest.mark.asyncio
    async def test_retrieve_disconnects_clients(
        self, mock_redis_client_instance, mock_mongo_client_instance, sample_context, sample_customer_id
    ):
        """Test that clients are disconnected after retrieval."""
        mock_redis_client_instance.get_context = AsyncMock(return_value=None)
        mock_mongo_client_instance.get_context = AsyncMock(return_value=sample_context)
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis_client_instance)
        ), patch(
            "src.geaicontext.tools.retrieve.get_mongo_client",
            AsyncMock(return_value=mock_mongo_client_instance)
        ):
            await retrieve_user_context(sample_customer_id)
            
            mock_redis_client_instance.disconnect.assert_called_once()
            mock_mongo_client_instance.disconnect.assert_called_once()

    @pytest.mark.asyncio
    async def test_retrieve_disconnects_on_error(
        self, mock_redis_client_instance, mock_mongo_client_instance, sample_customer_id
    ):
        """Test that clients are disconnected even on error."""
        mock_redis_client_instance.get_context = AsyncMock(return_value=None)
        mock_mongo_client_instance.get_context = AsyncMock(side_effect=PyMongoError("Error"))
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis_client_instance)
        ), patch(
            "src.geaicontext.tools.retrieve.get_mongo_client",
            AsyncMock(return_value=mock_mongo_client_instance)
        ):
            await retrieve_user_context(sample_customer_id)
            
            mock_redis_client_instance.disconnect.assert_called_once()
            mock_mongo_client_instance.disconnect.assert_called_once()
