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 get_metrics implementation
        self.mock_metrics = Mock(spec=SystemMetrics)
        self.mock_metrics.get_metrics.return_value = {
            "timestamp": "2024-10-31T12:00:00Z",
            "cpu": 10,
            "gpu": 20,
            "memory_used": 1000,
            "memory_total": 2000,
            "disk_used": 5000,
            "disk_total": 10000,
            "token_count": 100,
            "token_limit": 1000
        }
        
        # 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
        self.mock_llm.token_count.return_value = 100
        self.mock_llm.token_limit.return_value = 1000
        
        # 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,
            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

    async def _skip_initial_messages(self, ws):
        """Helper to skip past the initial connection messages"""
        for _ in range(4):  # Now 4 messages including auto approver config
            await ws.receive_json()

    @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'])

            # Verify initial auto approver config
            msg = await ws.receive_json()
            self.assertEqual(msg['type'], ServerMessage.AUTO_APPROVER_CONFIG)
            self.assertIn('config', msg)

    @unittest_run_loop
    async def test_approve_context(self):
        """Test context approval flow"""
        async with self.client.ws_connect('/ws') as ws:
            await self._skip_initial_messages(ws)
            
            await ws.send_json({
                'type': ClientMessage.APPROVE_CONTEXT
            })
            
            msg = await ws.receive_json()
            self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
            self.assertEqual(msg['state'], WebAgentState.INFERENCE.name)
            
            msg = await ws.receive_json()
            self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
            self.assertEqual(msg['response'], "")
            
            msg = await ws.receive_json()
            self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
            self.assertEqual(msg['response'], "<reasoning>test reasoning</reasoning>")
            
            msg = await ws.receive_json()
            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:
            await self._skip_initial_messages(ws)
                
            await ws.send_json({
                'type': ClientMessage.SEND_INPUT,
                'input': test_input
            })

            await asyncio.sleep(0.5)
            self.assertEqual(self.io_buffer.read(), test_input)

    @unittest_run_loop
    async def test_modify_response(self):
        """Test modifying response in RESPONSE_APPROVAL state"""
        async with self.client.ws_connect('/ws') as ws:
            await self._skip_initial_messages(ws)
            
            await ws.send_json({
                'type': ClientMessage.APPROVE_CONTEXT
            })
            
            # Skip to RESPONSE_APPROVAL state
            for _ in range(4):
                await ws.receive_json()
                
            modified_response = "<reasoning>modified test</reasoning>"
            await ws.send_json({
                'type': ClientMessage.MODIFY_RESPONSE,
                'response': modified_response
            })
            
            msg = await ws.receive_json()
            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"""
        async with self.client.ws_connect('/ws') as ws1, self.client.ws_connect('/ws') as ws2:
            await self._skip_initial_messages(ws1)
            await self._skip_initial_messages(ws2)
            
            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)

    @unittest_run_loop
    async def test_auto_approver_config(self):
        """Test setting and receiving auto approver config"""
        async with self.client.ws_connect('/ws') as ws:
            await self._skip_initial_messages(ws)
            
            test_config = {
                'context_enabled': True,
                'response_enabled': True,
                'context_timeout': 10.0,
                'response_timeout': 20.0
            }
            
            await ws.send_json({
                'type': ClientMessage.AUTO_APPROVER_CONFIG,
                'config': test_config
            })
            
            msg = await ws.receive_json()
            self.assertEqual(msg['type'], ServerMessage.AUTO_APPROVER_CONFIG)
            self.assertEqual(msg['config'], test_config)

            self.ws_manager.auto_approver.context_enabled = False
            self.ws_manager.auto_approver.response_enabled = False

    @unittest_run_loop
    async def test_auto_approver_broadcast(self):
        """Test broadcasting auto approver config changes to multiple clients"""
        async with self.client.ws_connect('/ws') as ws1, self.client.ws_connect('/ws') as ws2:
            await self._skip_initial_messages(ws1)
            await self._skip_initial_messages(ws2)
            
            test_config = {
                'context_enabled': True,
                'response_enabled': False,
                'context_timeout': 15.0,
                'response_timeout': 25.0
            }
            
            await ws1.send_json({
                'type': ClientMessage.AUTO_APPROVER_CONFIG,
                'config': test_config
            })
            
            msg1 = await ws1.receive_json()
            msg2 = await ws2.receive_json()
            self.assertEqual(msg1['type'], ServerMessage.AUTO_APPROVER_CONFIG)
            self.assertEqual(msg2['type'], ServerMessage.AUTO_APPROVER_CONFIG)
            self.assertEqual(msg1['config'], test_config)
            self.assertEqual(msg2['config'], test_config)

            self.ws_manager.auto_approver.context_enabled = False
            self.ws_manager.auto_approver.response_enabled = False

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