from aiohttp import WSMsgType
from aiohttp import WSMsgType
from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop
from pathlib import Path
from sia.web_agent import WebAgentState
from unittest.mock import patch
import asyncio
import asyncio
import json
import json
import shutil
import tempfile

from sia.web_agent import WebAgentState
from sia.web_io_buffer import WebIOBuffer
from sia.web_server import WebServer, ServerMessage, ClientMessage
from .mock_web_agent import MockWebAgent

class WebServerTest(AioHTTPTestCase):
    async def get_application(self):
        """Create application for testing."""
        self.io_buffer = WebIOBuffer()
        
        self.agent = MockWebAgent()

        # Create temporary directory for static files
        self.temp_dir = Path(tempfile.mkdtemp())
        self.static_dir = self.temp_dir / "static"
        self.static_dir.mkdir(exist_ok=True)
        
        # Create web server with temp static directory
        self.web_server = WebServer(
            self.agent,
            self.io_buffer,
            "localhost",
            8080,
            static_dir=self.static_dir
        )
        return self.web_server.app

    def tearDown(self):
        """Clean up temporary directory."""
        if hasattr(self, 'temp_dir'):
            shutil.rmtree(self.temp_dir)

    async def wait_for_messages(self, ws, count=1, timeout=5):
        """Helper to wait for multiple WebSocket messages."""
        messages = []
        try:
            for _ in range(count):
                msg = await ws.receive(timeout=timeout)
                if msg.type == WSMsgType.TEXT:
                    messages.append(json.loads(msg.data))
                elif msg.type == WSMsgType.BINARY:
                    messages.append(json.loads(msg.data.decode()))
                elif msg.type == WSMsgType.CLOSED:
                    print("WebSocket closed normally")
                    break
                elif msg.type == WSMsgType.ERROR:
                    self.fail(f"WebSocket error while waiting for messages: {msg.data}")
            return messages
        except asyncio.TimeoutError:
            self.fail(f"Timeout waiting for WebSocket message. Got {len(messages)} of {count}")
        except json.JSONDecodeError as e:
            self.fail(f"Invalid JSON in message: {e}")

    async def drain_messages(self, ws, timeout=0.1):
        """Helper to drain any pending messages."""
        try:
            while True:
                msg = await ws.receive(timeout=timeout)
                if msg.type == WSMsgType.CLOSED:
                    break  # Normal close, no need to fail
                elif msg.type == WSMsgType.ERROR:
                    self.fail(f"WebSocket error while draining messages: {msg.data}")
        except asyncio.TimeoutError:
            pass  # Expected when no more messages

    def setUp(self):
        """Set up test case with initial state."""
        super().setUp()
        self.io_buffer = None
        self.agent = None
        self.web_server = None

    async def get_application(self):
        """Create application for testing."""
        self.io_buffer = WebIOBuffer()
        self.agent = MockWebAgent()
        self.agent._current_state = WebAgentState.CONTEXT_APPROVAL  # Set initial state
        self.web_server = WebServer(self.agent, self.io_buffer, "localhost", 8080)
        return self.web_server.app

    @unittest_run_loop
    async def test_approve_context(self):
        """Test context approval flow."""
        async with self.client.ws_connect("/ws") as ws:
            # Drain initial messages
            await self.drain_messages(ws)
            await asyncio.sleep(0.1)  # Let connection settle
            
            # Set agent to CONTEXT_APPROVAL state
            self.agent._current_state = WebAgentState.CONTEXT_APPROVAL
            self.agent.context = "<test>context</test>"
            await asyncio.sleep(0.1)  # Let state changes propagate
            
            # Send context approval
            await ws.send_json({
                "type": ClientMessage.APPROVE_CONTEXT,
            })
            
            # Should receive 3 messages: INFERENCE state, response, RESPONSE_APPROVAL state
            messages = await self.wait_for_messages(ws, count=3)
            
            # Verify state change to INFERENCE
            self.assertEqual(messages[0]["type"], ServerMessage.STATE_CHANGE)
            self.assertEqual(messages[0]["data"], WebAgentState.INFERENCE.name)

            # Verify response
            self.assertEqual(messages[1]["type"], ServerMessage.RESPONSE_UPDATE)
            self.assertEqual(messages[1]["data"], "<reasoning>test</reasoning>")
            
            # Verify state change to RESPONSE_APPROVAL
            self.assertEqual(messages[2]["type"], ServerMessage.STATE_CHANGE)
            self.assertEqual(messages[2]["data"], WebAgentState.RESPONSE_APPROVAL.name)

    @unittest_run_loop        
    async def test_validation_error(self):
        """Test handling of validation errors."""
        async with self.client.ws_connect("/ws") as ws:
            # Drain initial messages
            await self.drain_messages(ws)
            await asyncio.sleep(0.1)  # Let connection settle
            
            # Set validation error and response
            self.agent._validation_error = "Invalid XML"
            await self.agent._set_response("<invalid>")
            await asyncio.sleep(0.1)  # Let changes propagate
            
            # Should receive response update with error
            messages = await self.wait_for_messages(ws, count=1)
            self.assertEqual(messages[0]["type"], ServerMessage.RESPONSE_UPDATE)
            self.assertEqual(messages[0]["validation_error"], "Invalid XML")
            
    @unittest_run_loop
    async def test_multiple_clients(self):
        """Test handling multiple client connections."""
        ws1 = await self.client.ws_connect("/ws")
        ws2 = await self.client.ws_connect("/ws")
        
        try:
            # Drain initial messages
            await self.drain_messages(ws1)
            await self.drain_messages(ws2)
            await asyncio.sleep(0.1)  # Let connections settle
            
            # Verify both clients are tracked
            self.assertEqual(len(self.web_server.clients), 2)
            
        finally:
            await ws1.close()
            await ws2.close()
            
    @unittest_run_loop
    async def test_client_disconnect(self):
        """Test proper cleanup on client disconnect."""
        ws = await self.client.ws_connect("/ws")
        await self.drain_messages(ws)
        await asyncio.sleep(0.1)  # Let connection settle
        
        self.assertEqual(len(self.web_server.clients), 1)
        
        await ws.close()
        await asyncio.sleep(0.1)  # Let cleanup complete
        
        self.assertEqual(len(self.web_server.clients), 0)
    
    @unittest_run_loop
    async def test_broadcast(self):
        """Test broadcasting to all clients."""
        ws1 = await self.client.ws_connect("/ws")
        ws2 = await self.client.ws_connect("/ws")
        
        try:
            # Drain initial messages
            await self.drain_messages(ws1)
            await self.drain_messages(ws2)
            await asyncio.sleep(0.1)  # Let connections settle
            
            # Broadcast test message
            test_msg = "test broadcast"
            await self.web_server.broadcast(ServerMessage.OUTPUT_UPDATE, test_msg)
            
            # Both clients should receive message
            msg1 = (await self.wait_for_messages(ws1, count=1))[0]
            msg2 = (await self.wait_for_messages(ws2, count=1))[0]
            
            self.assertEqual(msg1["type"], ServerMessage.OUTPUT_UPDATE)
            self.assertEqual(msg1["data"], test_msg)
            self.assertEqual(msg2["type"], ServerMessage.OUTPUT_UPDATE)
            self.assertEqual(msg2["data"], test_msg)
            
        finally:
            await ws1.close()
            await ws2.close()
            
    @unittest_run_loop
    async def test_send_input(self):
        """Test sending input to agent."""
        async with self.client.ws_connect("/ws") as ws:
            # Drain initial messages
            await self.drain_messages(ws)
            await asyncio.sleep(0.1)  # Let connection settle
            
            # Send input
            test_input = "test input"
            await ws.send_json({
                "type": ClientMessage.SEND_INPUT,
                "data": test_input
            })
            await asyncio.sleep(0.1)  # Let input process
            
            # Verify input added to buffer with newline
            self.assertEqual(self.io_buffer.read(), test_input + "\n")

    @unittest_run_loop
    async def test_invalid_json(self):
        """Test handling of invalid JSON messages."""
        async with self.client.ws_connect("/ws") as ws:
            # Drain initial messages
            await self.drain_messages(ws)
            await asyncio.sleep(0.1)  # Let connection settle
            
            # Send invalid JSON
            await ws.send_str("invalid json")
            await asyncio.sleep(0.1)  # Let error process
            
            # Connection should remain open
            self.assertEqual(len(self.web_server.clients), 1)