"""Tests for MCP Client Manager."""

import pytest
from unittest.mock import AsyncMock, patch
from mcp_forge.mcp.manager import MCPClientManager


@pytest.fixture
def mock_client1():
    """Create mock MCP client 1."""
    client = AsyncMock()
    client.name = "client1"
    client.is_connected.return_value = False
    client.connect = AsyncMock()
    client.disconnect = AsyncMock()
    client.list_tools = AsyncMock(return_value=["tool1", "tool2"])
    client.get_tool_schema = AsyncMock(return_value={
        "type": "object",
        "properties": {"param1": {"type": "string"}}
    })
    client.call_tool = AsyncMock(return_value={"result": "success"})
    return client


@pytest.fixture
def mock_client2():
    """Create mock MCP client 2."""
    client = AsyncMock()
    client.name = "client2"
    client.is_connected.return_value = False
    client.connect = AsyncMock()
    client.disconnect = AsyncMock()
    client.list_tools = AsyncMock(return_value=["tool3", "tool4"])
    client.get_tool_schema = AsyncMock(return_value={
        "type": "object",
        "properties": {"param2": {"type": "number"}}
    })
    client.call_tool = AsyncMock(return_value={"result": "success2"})
    return client


@pytest.fixture
def client_config():
    """Create test client configuration."""
    return {
        "client1": {
            "command": "python",
            "args": ["server1.py"],
            "env": {"KEY1": "value1"}
        },
        "client2": {
            "command": "python",
            "args": ["server2.py"],
            "env": {"KEY2": "value2"}
        }
    }


@pytest.mark.asyncio
async def test_initialize_clients_from_config(mock_client1, mock_client2, client_config):
    """Test initializing clients from configuration."""
    with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
        # Setup mock to return different clients for different configs
        mock_wrapper.side_effect = [mock_client1, mock_client2]
        
        manager = MCPClientManager(client_config)
        await manager.initialize()
        
        # Verify clients were created with correct config
        assert mock_wrapper.call_count == 2
        
        # Verify first client created with correct params
        call1 = mock_wrapper.call_args_list[0]
        assert call1[1]["name"] == "client1"
        assert call1[1]["command"] == "python"
        assert call1[1]["args"] == ["server1.py"]
        assert call1[1]["env"] == {"KEY1": "value1"}
        
        # Verify second client created with correct params
        call2 = mock_wrapper.call_args_list[1]
        assert call2[1]["name"] == "client2"
        assert call2[1]["command"] == "python"
        assert call2[1]["args"] == ["server2.py"]
        assert call2[1]["env"] == {"KEY2": "value2"}
        
        # Verify clients were connected
        mock_client1.connect.assert_called_once()
        mock_client2.connect.assert_called_once()
        
        # Verify tools were listed
        mock_client1.list_tools.assert_called_once()
        mock_client2.list_tools.assert_called_once()


@pytest.mark.asyncio
async def test_get_client_for_tool(mock_client1, mock_client2, client_config):
    """Test getting client that provides a specific tool."""
    with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
        mock_wrapper.side_effect = [mock_client1, mock_client2]
        
        manager = MCPClientManager(client_config)
        await manager.initialize()
        
        # Get client for tool1 (from client1)
        client = await manager.get_client_for_tool("tool1")
        assert client == mock_client1
        
        # Get client for tool3 (from client2)
        client = await manager.get_client_for_tool("tool3")
        assert client == mock_client2


@pytest.mark.asyncio
async def test_get_client_for_unknown_tool(mock_client1, mock_client2, client_config):
    """Test getting client for tool that doesn't exist."""
    with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
        mock_wrapper.side_effect = [mock_client1, mock_client2]
        
        manager = MCPClientManager(client_config)
        await manager.initialize()
        
        # Try to get client for non-existent tool
        with pytest.raises(KeyError, match="Tool 'unknown_tool' not found"):
            await manager.get_client_for_tool("unknown_tool")


@pytest.mark.asyncio
async def test_list_all_tools(mock_client1, mock_client2, client_config):
    """Test listing all tools across all clients."""
    with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
        mock_wrapper.side_effect = [mock_client1, mock_client2]
        
        manager = MCPClientManager(client_config)
        await manager.initialize()
        
        tools = await manager.list_all_tools()
        
        # Should have all tools from both clients
        assert set(tools) == {"tool1", "tool2", "tool3", "tool4"}


@pytest.mark.asyncio
async def test_detect_tool_name_collision():
    """Test detection of tool name collisions across clients."""
    # Create clients with overlapping tool names
    client1 = AsyncMock()
    client1.name = "client1"
    client1.connect = AsyncMock()
    client1.list_tools = AsyncMock(return_value=["tool1", "tool2"])
    
    client2 = AsyncMock()
    client2.name = "client2"
    client2.connect = AsyncMock()
    client2.list_tools = AsyncMock(return_value=["tool2", "tool3"])  # tool2 collision!
    
    config = {
        "client1": {"command": "python", "args": ["server1.py"]},
        "client2": {"command": "python", "args": ["server2.py"]}
    }
    
    with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
        mock_wrapper.side_effect = [client1, client2]
        
        manager = MCPClientManager(config)
        
        # Should raise ValueError about collision during initialization
        with pytest.raises(ValueError, match="Tool name collision.*tool2.*client1.*client2"):
            await manager.initialize()


@pytest.mark.asyncio
async def test_get_tool_schema(mock_client1, mock_client2, client_config):
    """Test getting tool schema via manager."""
    with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
        mock_wrapper.side_effect = [mock_client1, mock_client2]
        
        manager = MCPClientManager(client_config)
        await manager.initialize()
        
        # Get schema for tool1 (from client1)
        schema = await manager.get_tool_schema("tool1")
        assert schema == {"type": "object", "properties": {"param1": {"type": "string"}}}
        mock_client1.get_tool_schema.assert_called_once_with("tool1")
        
        # Get schema for tool3 (from client2)
        schema = await manager.get_tool_schema("tool3")
        assert schema == {"type": "object", "properties": {"param2": {"type": "number"}}}
        mock_client2.get_tool_schema.assert_called_once_with("tool3")


@pytest.mark.asyncio
async def test_call_tool(mock_client1, mock_client2, client_config):
    """Test calling tool via manager."""
    with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
        mock_wrapper.side_effect = [mock_client1, mock_client2]
        
        manager = MCPClientManager(client_config)
        await manager.initialize()
        
        # Call tool1 (from client1)
        result = await manager.call_tool("tool1", {"param1": "value"})
        assert result == {"result": "success"}
        mock_client1.call_tool.assert_called_once_with("tool1", {"param1": "value"})
        
        # Call tool3 (from client2)
        result = await manager.call_tool("tool3", {"param2": 42})
        assert result == {"result": "success2"}
        mock_client2.call_tool.assert_called_once_with("tool3", {"param2": 42})


@pytest.mark.asyncio
async def test_shutdown_all_clients(mock_client1, mock_client2, client_config):
    """Test shutting down all clients."""
    with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
        mock_wrapper.side_effect = [mock_client1, mock_client2]
        
        manager = MCPClientManager(client_config)
        await manager.initialize()
        
        # Shutdown all clients
        await manager.shutdown()
        
        # Verify both clients were disconnected
        mock_client1.disconnect.assert_called_once()
        mock_client2.disconnect.assert_called_once()


@pytest.mark.asyncio
async def test_manager_before_initialization():
    """Test that manager methods fail before initialization."""
    config = {"client1": {"command": "python", "args": ["server.py"]}}
    manager = MCPClientManager(config)
    
    # Should raise RuntimeError if not initialized
    with pytest.raises(RuntimeError, match="Manager not initialized"):
        await manager.list_all_tools()
    
    with pytest.raises(RuntimeError, match="Manager not initialized"):
        await manager.get_client_for_tool("tool1")
    
    with pytest.raises(RuntimeError, match="Manager not initialized"):
        await manager.get_tool_schema("tool1")
    
    with pytest.raises(RuntimeError, match="Manager not initialized"):
        await manager.call_tool("tool1", {})