"""Tests for context.py MCP server module."""

import pytest
from unittest.mock import AsyncMock, MagicMock, patch
import os

from src.geaicontext.context import (
    tool_wrapper,
    main,
    mcp,
)


class TestToolWrapper:
    """Test cases for tool_wrapper decorator."""

    @pytest.mark.asyncio
    async def test_wrapper_success(self):
        """Test wrapper with successful function execution."""
        async def mock_fn(customer_id: str, data: dict):
            return {"result": "success"}
        
        wrapped = tool_wrapper(mock_fn)
        result = await wrapped("CUST-123", data={"key": "value"})
        
        assert result == {"result": "success"}

    @pytest.mark.asyncio
    async def test_wrapper_success_with_args(self):
        """Test wrapper extracts customer_id from args."""
        async def mock_fn(customer_id: str):
            return customer_id
        
        wrapped = tool_wrapper(mock_fn)
        result = await wrapped("CUST-456")
        
        assert result == "CUST-456"

    @pytest.mark.asyncio
    async def test_wrapper_success_with_kwargs(self):
        """Test wrapper extracts customer_id from kwargs."""
        async def mock_fn(customer_id: str):
            return customer_id
        
        wrapped = tool_wrapper(mock_fn)
        result = await wrapped(customer_id="CUST-789")
        
        assert result == "CUST-789"

    @pytest.mark.asyncio
    async def test_wrapper_no_customer_id(self):
        """Test wrapper handles missing customer_id."""
        async def mock_fn():
            return "no_id"
        
        wrapped = tool_wrapper(mock_fn)
        result = await wrapped()
        
        assert result == "no_id"

    @pytest.mark.asyncio
    async def test_wrapper_exception_handling(self):
        """Test wrapper handles exceptions gracefully."""
        async def mock_fn(customer_id: str):
            raise ValueError("Test error")
        
        wrapped = tool_wrapper(mock_fn)
        result = await wrapped("CUST-123")
        
        assert result["success"] is False
        assert result["error_type"] == "ValueError"
        assert "Test error" in result["error_message"]
        assert "user_message" in result

    @pytest.mark.asyncio
    async def test_wrapper_exception_with_different_types(self):
        """Test wrapper handles different exception types."""
        async def mock_fn(customer_id: str):
            raise RuntimeError("Runtime error")
        
        wrapped = tool_wrapper(mock_fn)
        result = await wrapped("CUST-123")
        
        assert result["error_type"] == "RuntimeError"


class TestSaveContextLogic:
    """Test cases for save_context logic (testing the underlying implementation)."""

    @pytest.mark.asyncio
    async def test_save_context_success(self, sample_context, sample_customer_id):
        """Test successful context save logic."""
        from src.geaicontext.tools.save import save_user_context
        
        mock_redis = MagicMock()
        mock_redis.connect = AsyncMock()
        mock_redis.disconnect = AsyncMock()
        mock_redis.set_context = AsyncMock(return_value=True)
        
        mock_mongo = MagicMock()
        mock_mongo.connect = AsyncMock()
        mock_mongo.disconnect = AsyncMock()
        mock_mongo.save_context = AsyncMock(return_value=True)
        
        with patch(
            "src.geaicontext.tools.save.get_redis_client",
            AsyncMock(return_value=mock_redis)
        ), patch(
            "src.geaicontext.tools.save.get_mongo_client",
            AsyncMock(return_value=mock_mongo)
        ):
            success, message = await save_user_context(sample_customer_id, sample_context)
            
            assert success is True
            assert sample_customer_id in message

    @pytest.mark.asyncio
    async def test_save_context_validation_failure(self, sample_context, sample_customer_id):
        """Test save_context with validation failure."""
        from src.geaicontext.tools.save import save_user_context
        
        mock_redis = MagicMock()
        mock_redis.connect = AsyncMock()
        mock_redis.disconnect = AsyncMock()
        mock_redis.set_context = AsyncMock(return_value=False)
        
        with patch(
            "src.geaicontext.tools.save.get_redis_client",
            AsyncMock(return_value=mock_redis)
        ):
            success, message = await save_user_context(sample_customer_id, sample_context)
            
            assert success is False
            assert "Failed to write to Redis cache" in message


class TestRetrieveContextLogic:
    """Test cases for retrieve_context logic (testing the underlying implementation)."""

    @pytest.mark.asyncio
    async def test_retrieve_context_success(self, sample_context, sample_customer_id):
        """Test successful context retrieval logic."""
        from src.geaicontext.tools.retrieve import retrieve_user_context
        
        mock_redis = MagicMock()
        mock_redis.connect = AsyncMock()
        mock_redis.disconnect = AsyncMock()
        mock_redis.get_context = AsyncMock(return_value=sample_context)
        
        mock_mongo = MagicMock()
        mock_mongo.connect = AsyncMock()
        mock_mongo.disconnect = AsyncMock()
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis)
        ), patch(
            "src.geaicontext.tools.retrieve.get_mongo_client",
            AsyncMock(return_value=mock_mongo)
        ):
            result = await retrieve_user_context(sample_customer_id)
            
            assert result is not None
            assert result == sample_context

    @pytest.mark.asyncio
    async def test_retrieve_context_not_found(self, sample_customer_id):
        """Test retrieve_context when context not found."""
        from src.geaicontext.tools.retrieve import retrieve_user_context
        
        mock_redis = MagicMock()
        mock_redis.connect = AsyncMock()
        mock_redis.disconnect = AsyncMock()
        mock_redis.get_context = AsyncMock(return_value=None)
        
        mock_mongo = MagicMock()
        mock_mongo.connect = AsyncMock()
        mock_mongo.disconnect = AsyncMock()
        mock_mongo.get_context = AsyncMock(return_value=None)
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis)
        ), patch(
            "src.geaicontext.tools.retrieve.get_mongo_client",
            AsyncMock(return_value=mock_mongo)
        ):
            result = await retrieve_user_context(sample_customer_id)
            
            assert result is None


class TestServerConfiguration:
    """Test cases for server configuration."""

    def test_server_name_default(self):
        """Test default server name."""
        # The server_name is set at module load time
        # We can verify the mcp instance has a name
        assert mcp.name is not None

    def test_server_name_from_env(self):
        """Test server name from environment variable."""
        with patch.dict(os.environ, {"SERVER_NAME": "Custom MCP Server"}):
            # Re-import to get new value
            # Note: This tests the pattern, actual reload would be needed
            assert os.getenv("SERVER_NAME") == "Custom MCP Server"


class TestMain:
    """Test cases for main function."""

    def test_main_calls_mcp_run(self):
        """Test that main() calls mcp.run()."""
        with patch.object(mcp, "run") as mock_run:
            main()
            mock_run.assert_called_once()


class TestMCPToolRegistration:
    """Test cases for MCP tool registration."""

    def test_save_context_tool_exists(self):
        """Test that save_context is registered as an MCP tool."""
        # The @mcp.tool() decorator wraps the function in a FunctionTool
        # We verify the tool is registered by checking the mcp instance
        from src.geaicontext.context import save_context
        # FunctionTool has a name attribute
        assert hasattr(save_context, 'name')
        assert save_context.name == 'save_context'

    def test_retrieve_context_tool_exists(self):
        """Test that retrieve_context is registered as an MCP tool."""
        from src.geaicontext.context import retrieve_context
        assert hasattr(retrieve_context, 'name')
        assert retrieve_context.name == 'retrieve_context'

    def test_save_context_has_description(self):
        """Test that save_context tool has a description."""
        from src.geaicontext.context import save_context
        assert hasattr(save_context, 'description')
        assert "Save user context" in save_context.description

    def test_retrieve_context_has_description(self):
        """Test that retrieve_context tool has a description."""
        from src.geaicontext.context import retrieve_context
        assert hasattr(retrieve_context, 'description')
        assert "Retrieve user context" in retrieve_context.description


class TestContextIntegration:
    """Integration tests for context module functions."""

    @pytest.mark.asyncio
    async def test_tool_wrapper_with_save_user_context(self, sample_context, sample_customer_id):
        """Test tool_wrapper integration with save_user_context."""
        from src.geaicontext.tools.save import save_user_context
        
        mock_redis = MagicMock()
        mock_redis.connect = AsyncMock()
        mock_redis.disconnect = AsyncMock()
        mock_redis.set_context = AsyncMock(return_value=True)
        
        mock_mongo = MagicMock()
        mock_mongo.connect = AsyncMock()
        mock_mongo.disconnect = AsyncMock()
        mock_mongo.save_context = AsyncMock(return_value=True)
        
        with patch(
            "src.geaicontext.tools.save.get_redis_client",
            AsyncMock(return_value=mock_redis)
        ), patch(
            "src.geaicontext.tools.save.get_mongo_client",
            AsyncMock(return_value=mock_mongo)
        ):
            wrapped_fn = tool_wrapper(save_user_context)
            result = await wrapped_fn(sample_customer_id, sample_context)
            
            # Result should be a tuple (success, message)
            assert isinstance(result, tuple)
            assert result[0] is True

    @pytest.mark.asyncio
    async def test_tool_wrapper_with_retrieve_user_context(self, sample_context, sample_customer_id):
        """Test tool_wrapper integration with retrieve_user_context."""
        from src.geaicontext.tools.retrieve import retrieve_user_context
        
        mock_redis = MagicMock()
        mock_redis.connect = AsyncMock()
        mock_redis.disconnect = AsyncMock()
        mock_redis.get_context = AsyncMock(return_value=sample_context)
        
        mock_mongo = MagicMock()
        mock_mongo.connect = AsyncMock()
        mock_mongo.disconnect = AsyncMock()
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis)
        ), patch(
            "src.geaicontext.tools.retrieve.get_mongo_client",
            AsyncMock(return_value=mock_mongo)
        ):
            wrapped_fn = tool_wrapper(retrieve_user_context)
            result = await wrapped_fn(sample_customer_id)
            
            assert result == sample_context

    @pytest.mark.asyncio
    async def test_tool_wrapper_error_handling_integration(self, sample_customer_id):
        """Test tool_wrapper error handling with real function."""
        from src.geaicontext.tools.retrieve import retrieve_user_context
        from redis.exceptions import RedisError
        
        mock_redis = MagicMock()
        mock_redis.connect = AsyncMock(side_effect=RedisError("Connection failed"))
        
        with patch(
            "src.geaicontext.tools.retrieve.get_redis_client",
            AsyncMock(return_value=mock_redis)
        ):
            wrapped_fn = tool_wrapper(retrieve_user_context)
            result = await wrapped_fn(sample_customer_id)
            
            # The wrapper should catch the exception and return error dict
            assert isinstance(result, dict)
            assert result["success"] is False
            assert "error_type" in result
