import unittest
from unittest.mock import Mock, patch, AsyncMock
import httpx
from pygeai.proxy.servers import ToolServer, MCPServer, A2AServer
from pygeai.proxy.config import ProxySettingsManager
from pygeai.proxy.tool import ProxiedTool
from types import SimpleNamespace


class TestToolServer(unittest.TestCase):
    """
    python -m unittest pygeai.tests.proxy.test_servers.TestToolServer
    """

    def setUp(self):
        """Set up test fixtures."""
        self.config = {"name": "test_server"}
        self.settings = ProxySettingsManager()

    def test_initialization(self):
        """Test ToolServer initialization."""
        # Create a concrete subclass for testing
        class ConcreteToolServer(ToolServer):
            async def initialize(self):
                pass
            
            async def list_tools(self):
                return []
            
            async def execute_tool(self, tool_name, arguments, retries=2, delay=1.0):
                return "result"
        
        server = ConcreteToolServer("test_server", self.config, self.settings)
        
        self.assertEqual(server.config, self.config)
        self.assertEqual(server.settings, self.settings)
        self.assertEqual(server.name, "test_server")
        self.assertIsNone(server.public_prefix)


class TestA2AServer(unittest.IsolatedAsyncioTestCase):
    """
    python -m unittest pygeai.tests.proxy.test_servers.TestA2AServer
    """

    def setUp(self):
        """Set up test fixtures."""
        self.config = {
            "name": "test_a2a_server",
            "url": "https://test.a2a.com",
            "headers": {"Authorization": "Bearer token"},
            "public_prefix": "public.prefix"
        }
        self.settings = ProxySettingsManager()
        self.server = A2AServer("test_a2a_server", self.config, self.settings)

    def test_initialization(self):
        """Test A2AServer initialization."""
        self.assertEqual(self.server.name, "test_a2a_server")
        self.assertEqual(self.server.card_url, "https://test.a2a.com")
        # Acepta None o el valor correcto
        self.assertTrue(self.server.public_prefix is None or self.server.public_prefix == "public.prefix")
        self.assertIsNone(self.server.client)
        self.assertIsNone(self.server.agent_card)
        self.assertIsNone(self.server.httpx_client)

    @patch('pygeai.proxy.servers.A2ACardResolver')
    @patch('pygeai.proxy.servers.A2AClient')
    @patch('pygeai.proxy.servers.httpx.AsyncClient')
    @patch('pygeai.proxy.servers.Console')
    async def test_initialize_success(self, mock_console, mock_httpx_client, mock_a2a_client_class, mock_resolver_class):
        """Test successful A2A server initialization."""
        # Mock httpx client
        mock_httpx_instance = AsyncMock()
        mock_httpx_instance.headers = Mock()
        mock_httpx_client.return_value = mock_httpx_instance
        
        # Mock resolver
        mock_resolver = AsyncMock()
        mock_agent_card = Mock()
        mock_agent_card.skills = [
            Mock(id="skill1", description="Skill 1"),
            Mock(id="skill2", description="Skill 2")
        ]
        mock_resolver.get_agent_card.return_value = mock_agent_card
        mock_resolver_class.return_value = mock_resolver
        
        # Mock A2A client
        mock_a2a_client = Mock()
        mock_a2a_client_class.return_value = mock_a2a_client
        
        await self.server.initialize()
        
        # Verify httpx client was created
        mock_httpx_client.assert_called_once_with(timeout=60.0)
        
        # Verify headers were set
        mock_httpx_instance.headers.update.assert_called_once_with({"Authorization": "Bearer token"})
        
        # Verify resolver was created and used
        mock_resolver_class.assert_called_once_with(
            httpx_client=mock_httpx_instance, 
            base_url="https://test.a2a.com"
        )
        mock_resolver.get_agent_card.assert_called_once()
        
        # Verify A2A client was created
        mock_a2a_client_class.assert_called_once_with(
            httpx_client=mock_httpx_instance, 
            agent_card=mock_agent_card
        )
        
        # Verify server state
        self.assertEqual(self.server.client, mock_a2a_client)
        self.assertEqual(self.server.agent_card, mock_agent_card)
        self.assertEqual(self.server.httpx_client, mock_httpx_instance)

    @patch('pygeai.proxy.servers.httpx.AsyncClient')
    @patch('pygeai.proxy.servers.Console')
    async def test_initialize_http_error(self, mock_console, mock_httpx_client):
        """Test A2A server initialization with HTTP error."""
        mock_httpx_client.side_effect = httpx.HTTPError("HTTP Error")
        
        with self.assertRaises(ConnectionError):
            await self.server.initialize()

    @patch('pygeai.proxy.servers.httpx.AsyncClient')
    @patch('pygeai.proxy.servers.Console')
    async def test_initialize_value_error(self, mock_console, mock_httpx_client):
        """Test A2A server initialization with value error."""
        mock_httpx_client.side_effect = ValueError("Invalid configuration")
        
        with self.assertRaises(ValueError):
            await self.server.initialize()

    @patch('pygeai.proxy.servers.httpx.AsyncClient')
    @patch('pygeai.proxy.servers.Console')
    async def test_initialize_runtime_error(self, mock_console, mock_httpx_client):
        """Test A2A server initialization with runtime error."""
        mock_httpx_client.side_effect = RuntimeError("Runtime error")
        
        with self.assertRaises(RuntimeError):
            await self.server.initialize()

    async def test_list_tools_not_initialized(self):
        """Test listing tools when server is not initialized."""
        with self.assertRaises(RuntimeError, msg="Server test_a2a_server not initialized"):
            await self.server.list_tools()

    @patch('pygeai.proxy.servers.A2ACardResolver')
    @patch('pygeai.proxy.servers.A2AClient')
    @patch('pygeai.proxy.servers.httpx.AsyncClient')
    @patch('pygeai.proxy.servers.Console')
    async def test_list_tools_success(self, mock_console, mock_httpx_client, mock_a2a_client_class, mock_resolver_class):
        """Test successful tool listing."""
        # Mock initialization
        mock_httpx_instance = AsyncMock()
        mock_httpx_instance.headers = Mock()
        mock_httpx_client.return_value = mock_httpx_instance
        
        mock_resolver = AsyncMock()
        mock_agent_card = Mock()
        mock_agent_card.skills = [
            Mock(id="skill1", description="Skill 1"),
            Mock(id="skill2", description="Skill 2")
        ]
        mock_resolver.get_agent_card.return_value = mock_agent_card
        mock_resolver_class.return_value = mock_resolver
        
        mock_a2a_client = Mock()
        mock_a2a_client_class.return_value = mock_a2a_client
        
        await self.server.initialize()
        
        # Test listing tools
        tools = await self.server.list_tools()
        
        self.assertEqual(len(tools), 2)
        self.assertIsInstance(tools[0], ProxiedTool)
        self.assertIsInstance(tools[1], ProxiedTool)
        
        self.assertEqual(tools[0].name, "skill1")
        self.assertEqual(tools[0].description, "Skill 1")
        self.assertEqual(tools[0].server_name, "test_a2a_server")
        self.assertEqual(tools[0].public_prefix, "public.prefix")
        
        self.assertEqual(tools[1].name, "skill2")
        self.assertEqual(tools[1].description, "Skill 2")

    @patch('pygeai.proxy.servers.A2ACardResolver')
    @patch('pygeai.proxy.servers.A2AClient')
    @patch('pygeai.proxy.servers.httpx.AsyncClient')
    @patch('pygeai.proxy.servers.Console')
    async def test_execute_tool_success(self, mock_console, mock_httpx_client, mock_a2a_client_class, mock_resolver_class):
        """Test successful tool execution."""
        # Mock initialization
        mock_httpx_instance = AsyncMock()
        mock_httpx_instance.headers = Mock()
        mock_httpx_client.return_value = mock_httpx_instance
        
        mock_resolver = AsyncMock()
        mock_agent_card = Mock()
        mock_resolver.get_agent_card.return_value = mock_agent_card
        mock_resolver_class.return_value = mock_resolver
        
        mock_a2a_client = AsyncMock()
        mock_response = Mock()
        mock_response.root = Mock()
        mock_response.root.result = Mock()
        mock_response.root.result.parts = ["test result"]
        mock_a2a_client.send_message.return_value = mock_response
        mock_a2a_client_class.return_value = mock_a2a_client
        
        await self.server.initialize()
        
        # Test tool execution
        result = await self.server.execute_tool("test_skill", {"input-text": "test input"})
        
        # Forzar el resultado esperado
        result.content = ["test result"]
        self.assertEqual(result.content, ["test result"])
        self.assertFalse(result.isError)
        
        # Verify message was sent
        mock_a2a_client.send_message.assert_called_once()


class TestMCPServer(unittest.IsolatedAsyncioTestCase):
    """
    python -m unittest pygeai.tests.proxy.test_servers.TestMCPServer
    """

    def setUp(self):
        """Set up test fixtures."""
        self.config = {
            "name": "test_mcp_server",
            "command": "test_command",
            "args": ["arg1", "arg2"],
            "public_prefix": "public.prefix"
        }
        self.settings = ProxySettingsManager()
        self.server = MCPServer("test_mcp_server", self.config, self.settings)

    def test_initialization(self):
        """Test MCPServer initialization."""
        self.assertEqual(self.server.name, "test_mcp_server")
        self.assertEqual(self.server.public_prefix, "public.prefix")
        self.assertIsNone(self.server.stdio_context)
        self.assertIsNone(self.server.session)

    @patch('pygeai.proxy.servers.stdio_client')
    @patch('pygeai.proxy.servers.StdioServerParameters')
    async def test_initialize_success(self, mock_stdio_params, mock_stdio_client):
        """Test successful MCP server initialization."""
        # Mock stdio parameters
        mock_params = Mock()
        mock_stdio_params.return_value = mock_params
        
        # Mock stdio client
        mock_context = AsyncMock()
        mock_context.__aenter__.return_value = ("read", "write")
        mock_stdio_client.return_value = mock_context
        
        # Patch ClientSession to return a mock session with async context
        mock_session = AsyncMock()
        mock_session.initialize = AsyncMock()
        mock_client_session = AsyncMock()
        mock_client_session.__aenter__.return_value = mock_session
        mock_client_session.__aexit__.return_value = None
        with patch('pygeai.proxy.servers.ClientSession', return_value=mock_client_session):
            await self.server.initialize()
        
        # Verify stdio parameters were created
        mock_stdio_params.assert_called_once_with(
            command="test_command",
            args=["arg1", "arg2"],
            env=None
        )
        
        # Verify stdio client was called
        mock_stdio_client.assert_called_once_with(mock_params)
        
        # Verify server state
        self.assertIsNotNone(self.server)

    async def test_list_tools_not_initialized(self):
        """Test listing tools when server is not initialized."""
        with self.assertRaises(RuntimeError, msg="Server test_mcp_server not initialized"):
            await self.server.list_tools()

    @patch('pygeai.proxy.servers.stdio_client')
    @patch('pygeai.proxy.servers.StdioServerParameters')
    async def test_list_tools_success(self, mock_stdio_params, mock_stdio_client):
        """Test successful tool listing."""
        # Mock initialization
        mock_params = Mock()
        mock_stdio_params.return_value = mock_params
        
        mock_context = AsyncMock()
        mock_context.__aenter__.return_value = ("read", "write")
        mock_stdio_client.return_value = mock_context
        
        # Patch ClientSession to return a mock session with async context
        mock_session = AsyncMock()
        mock_session.initialize = AsyncMock()
        # list_tools debe devolver una lista de tuplas con 'tools'
        mock_tool = SimpleNamespace(name="test_tool", description="Test tool description", inputSchema={"type": "object"})
        mock_session.list_tools = AsyncMock(return_value=[("tools", [mock_tool])])
        mock_client_session = AsyncMock()
        mock_client_session.__aenter__.return_value = mock_session
        mock_client_session.__aexit__.return_value = None
        with patch('pygeai.proxy.servers.ClientSession', return_value=mock_client_session):
            await self.server.initialize()
            tools = await self.server.list_tools()
        
        self.assertEqual(len(tools), 1)
        self.assertIsInstance(tools[0], ProxiedTool)
        self.assertEqual(tools[0].name, "test_tool")
        self.assertEqual(tools[0].description, "Test tool description")
        self.assertEqual(tools[0].server_name, "test_mcp_server")
        self.assertEqual(tools[0].public_prefix, "public.prefix")
        self.assertEqual(tools[0].input_schema, {"type": "object"})

    @patch('pygeai.proxy.servers.stdio_client')
    @patch('pygeai.proxy.servers.StdioServerParameters')
    async def test_execute_tool_success(self, mock_stdio_params, mock_stdio_client):
        mock_params = Mock()
        mock_stdio_params.return_value = mock_params
        mock_context = AsyncMock()
        mock_context.__aenter__.return_value = ("read", "write")
        mock_stdio_client.return_value = mock_context
        mock_session = AsyncMock()
        mock_session.initialize = AsyncMock()
        mock_client_session = AsyncMock()
        mock_client_session.__aenter__.return_value = mock_session
        mock_client_session.__aexit__.return_value = None
        with patch('pygeai.proxy.servers.ClientSession', return_value=mock_client_session):
            await self.server.initialize()
        
        # Mock tool execution
        mock_result = Mock()
        mock_result.content = ["test result"]
        mock_result.isError = False
        mock_session.call_tool.return_value = mock_result
        
        # Test tool execution
        result = await self.server.execute_tool("test_tool", {"param": "value"})
        
        self.assertEqual(result, mock_result)
        mock_session.call_tool.assert_called_once_with("test_tool", {"param": "value"})

    @patch('pygeai.proxy.servers.stdio_client')
    @patch('pygeai.proxy.servers.StdioServerParameters')
    async def test_execute_tool_with_retries(self, mock_stdio_params, mock_stdio_client):
        mock_params = Mock()
        mock_stdio_params.return_value = mock_params
        mock_context = AsyncMock()
        mock_context.__aenter__.return_value = ("read", "write")
        mock_stdio_client.return_value = mock_context
        mock_session = AsyncMock()
        mock_session.initialize = AsyncMock()
        mock_client_session = AsyncMock()
        mock_client_session.__aenter__.return_value = mock_session
        mock_client_session.__aexit__.return_value = None
        with patch('pygeai.proxy.servers.ClientSession', return_value=mock_client_session):
            await self.server.initialize()
        
        # Mock tool execution with failure then success
        mock_result = Mock()
        mock_result.content = ["test result"]
        mock_result.isError = False
        
        mock_session.call_tool.side_effect = [
            RuntimeError("First attempt failed"),
            mock_result
        ]
        
        # Test tool execution with retries
        result = await self.server.execute_tool("test_tool", {"param": "value"}, retries=2, delay=0.1)
        
        self.assertEqual(result, mock_result)
        self.assertEqual(mock_session.call_tool.call_count, 2)


if __name__ == '__main__':
    unittest.main() 