from aiohttp import web
from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop
from aiohttp.web import AppRunner
from unittest.mock import Mock, patch, call
import asyncio
import unittest
import xml.etree.ElementTree as ET

from sia.web_agent import WebAgent, WebAgentState
from sia.web_io_buffer import WebIOBuffer
from sia.web_socket_manager import WebSocketManager, ClientMessage, ServerMessage
from sia.working_memory import WorkingMemory
from sia.system_metrics import SystemMetrics
from sia.llm_engine import LlmEngine
from sia.xml_validator import XMLValidator
from sia.response_parser import ResponseParser

class WebSocketManagerTest(AioHTTPTestCase):
    async def get_application(self):
        """Create test application with WebSocket handler"""
        # Create real components where beneficial for testing
        self.io_buffer = WebIOBuffer()
        self.working_memory = WorkingMemory()
        
        # Create mock metrics with proper generate_context implementation
        self.mock_metrics = Mock(spec=SystemMetrics)
        metrics_elem = ET.Element("context")
        metrics_elem.attrib = {
            "time": "2024-10-31T12:00:00Z",
            "cpu": "10",
            "gpu": "20", 
            "memory_used": "1000",
            "memory_total": "2000",
            "disk_used": "5000",
            "disk_total": "10000",
            "context": "0",
            "stdin": "0"
        }
        self.mock_metrics.generate_context.return_value = metrics_elem
        
        # Create minimal mocks for other components
        self.mock_llm = Mock(spec=LlmEngine)
        def mock_infer(prompt: str, context: str):
            yield "<reasoning>test reasoning</reasoning>"
        self.mock_llm.infer.side_effect = mock_infer
        self.mock_validator = Mock(spec=XMLValidator)
        self.mock_validator.validate.return_value = None
        
        # Create parser with real IO buffer
        self.parser = ResponseParser(self.io_buffer)

        # Create agent with mix of real/mock components
        self.agent = WebAgent(
            system_prompt="test prompt",
            action_schema="test schema",
            working_memory=self.working_memory,
            system_metrics=self.mock_metrics,
            llm=self.mock_llm,
            validator=self.mock_validator,
            parser=self.parser
        )
        
        # Create WebSocketManager
        self.ws_manager = await WebSocketManager.create(self.agent, self.io_buffer)
        
        # Create application
        app = web.Application()
        app.router.add_get('/ws', self.ws_manager.handle_websocket)
        return app

    @unittest_run_loop
    async def test_initial_connection(self):
        """Test initial WebSocket connection and state messages"""
        async with self.client.ws_connect('/ws') as ws:
            # Verify initial state message
            msg = await ws.receive_json()
            self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
            self.assertEqual(msg['state'], WebAgentState.CONTEXT_APPROVAL.name)
            
            # Verify initial context message
            msg = await ws.receive_json()
            self.assertEqual(msg['type'], ServerMessage.CONTEXT_UPDATE)
            self.assertIsNotNone(msg['context'])
            
            # Verify initial response message
            msg = await ws.receive_json()
            self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
            self.assertEqual(msg['response'], '')
            self.assertIsNone(msg['validation_error'])

    @unittest_run_loop
    async def test_approve_context(self):
        """Test context approval flow"""
        async with self.client.ws_connect('/ws') as ws:
            for _ in range(3):
                msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
            await ws.send_json({
                'type': ClientMessage.APPROVE_CONTEXT
            })
            msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
            self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
            self.assertEqual(msg['state'], WebAgentState.INFERENCE.name)
            msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
            self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
            self.assertEqual(msg['response'], "")
            msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
            self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
            self.assertEqual(msg['response'], "<reasoning>test reasoning</reasoning>")
            msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
            self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
            self.assertEqual(msg['state'], WebAgentState.RESPONSE_APPROVAL.name)

    @unittest_run_loop
    async def test_send_input(self):
        """Test sending input through WebSocket"""
        test_input = "test input"
        
        async with self.client.ws_connect('/ws') as ws:
            for _ in range(3):
                await ws.receive_json()
                
            # Send input
            await ws.send_json({
                'type': ClientMessage.SEND_INPUT,
                'input': test_input
            })

            # Delay for processing
            await asyncio.sleep(0.5)
            
            # Verify input was added to buffer
            self.assertEqual(self.io_buffer.read(), test_input)

    @unittest_run_loop
    async def test_modify_response(self):
        """Test modifying response in RESPONSE_APPROVAL state"""
        # Set agent to RESPONSE_APPROVAL state
        async with self.client.ws_connect('/ws') as ws:
            for _ in range(3):
                msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
            await ws.send_json({
                'type': ClientMessage.APPROVE_CONTEXT
            })
            for _ in range(4):
                msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
            modified_response = "<reasoning>modified test</reasoning>"
            await ws.send_json({
                'type': ClientMessage.MODIFY_RESPONSE,
                'response': modified_response
            })
            msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
            self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
            self.assertEqual(msg['response'], modified_response)

    @unittest_run_loop
    async def test_multiple_clients(self):
        """Test broadcasting to multiple WebSocket clients"""
        # Connect first client
        async with self.client.ws_connect('/ws') as ws1, self.client.ws_connect('/ws') as ws2:
            for _ in range(3):
                await asyncio.wait_for(ws1.receive_json(), timeout=0.5)
                await asyncio.wait_for(ws2.receive_json(), timeout=0.5)
            await ws1.send_json({
                'type': ClientMessage.APPROVE_CONTEXT
            })
            msg1 = await ws1.receive_json()
            msg2 = await ws2.receive_json()
            self.assertEqual(msg1['type'], ServerMessage.STATE_CHANGE)
            self.assertEqual(msg1['state'], WebAgentState.INFERENCE.name)
            self.assertEqual(msg2['type'], ServerMessage.STATE_CHANGE)
            self.assertEqual(msg2['state'], WebAgentState.INFERENCE.name)

    def tearDown(self):
        """Clean up resources"""
        super().tearDown()
        self.io_buffer = None
        self.working_memory = None
        self.agent = None
        self.ws_manager = None