import asyncio
import unittest
from datetime import datetime
from unittest.mock import Mock, MagicMock, patch, call
import xml.etree.ElementTree as ET
from typing import Iterator

from sia.web_io_buffer import WebIOBuffer
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
from sia.web_agent import WebAgent, WebAgentState
from sia.command import Command
from sia.command_result import CommandResult
from sia.reasoning_entry import ReasoningEntry
from sia.parse_error_entry import ParseErrorEntry
from sia.entry import Entry

class MockCommand(Command):
    """Mock command for testing execution flow."""
    def __init__(self, result: CommandResult):
        self.result = result
        self.executed = False
        
    def execute(self, memory: WorkingMemory) -> CommandResult:
        self.executed = True
        return self.result

class WebAgentTest(unittest.TestCase):
    def setUp(self):
        """Set up test cases with mocked components."""
        # Create mocks
        self.mock_llm = Mock(spec=LlmEngine)
        self.mock_metrics = Mock(spec=SystemMetrics)
        self.mock_validator = Mock(spec=XMLValidator)
        self.io_buffer = WebIOBuffer()
        self.working_memory = WorkingMemory()
        
        # Mock validator to pass by default
        self.mock_validator.validate.return_value = None
        
        # Create parser with IO buffer
        self.parser = ResponseParser(self.io_buffer)
        
        # Mock metrics context generator
        def generate_context(context_size):
            metrics_elem = ET.Element("context")
            metrics_elem.set("time", "2024-10-31T12:00:00Z")
            metrics_elem.set("cpu", "10")
            metrics_elem.set("gpu", "20")
            metrics_elem.set("memory_used", "1000")
            metrics_elem.set("memory_total", "2000")
            metrics_elem.set("disk_used", "5000")
            metrics_elem.set("disk_total", "10000")
            metrics_elem.set("context", str(int(context_size * 100)))
            metrics_elem.set("stdin", "0")
            return metrics_elem
            
        self.mock_metrics.generate_context.side_effect = generate_context
        
        # Mock LLM response generator
        def mock_infer(prompt: str, context: str):
            yield "<reasoning>test reasoning</reasoning>"
            
        self.mock_llm.infer.side_effect = mock_infer
        
        # Create agent with all 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
        )
        
        # Set timestamp for entries
        self.test_timestamp = datetime(2024, 10, 31, 12, 0, 0)
        
        # Handler tracking
        self.state_changes = []
        self.response_changes = []
        
    def state_change_handler(self, state: WebAgentState):
        """Track state changes for verification."""
        self.state_changes.append(state)
        
    def response_change_handler(self, response: str):
        """Track response changes for verification."""
        self.response_changes.append(response)
        
    def test_initialization(self):
        """Test initial state of web agent."""
        self.assertEqual(self.agent.current_state, WebAgentState.CONTEXT_APPROVAL)
        self.assertEqual(self.agent.response, "")
        self.assertIsNotNone(self.agent.context)
        self.assertIsNone(self.agent._validation_error)
        self.assertEqual(len(self.agent._state_change_handlers), 0)
        self.assertEqual(len(self.agent._response_change_handlers), 0)
        
    def test_handler_registration(self):
        """Test adding state and response change handlers."""
        # Add handlers
        self.agent.add_state_change_handler(self.state_change_handler)
        self.agent.add_response_change_handler(self.response_change_handler)
        
        # Verify handlers added
        self.assertEqual(len(self.agent._state_change_handlers), 1)
        self.assertEqual(len(self.agent._response_change_handlers), 1)
        
        # Adding same handler twice should not duplicate
        self.agent.add_state_change_handler(self.state_change_handler)
        self.agent.add_response_change_handler(self.response_change_handler)
        
        self.assertEqual(len(self.agent._state_change_handlers), 1)
        self.assertEqual(len(self.agent._response_change_handlers), 1)
        
    def test_approve_context_state_error(self):
        """Test approving context in wrong state."""
        # Force agent into UPDATE state
        self.agent._current_state = WebAgentState.UPDATE
        
        with self.assertRaises(Exception) as context:
            asyncio.run(self.agent.approve_context())
        self.assertIn("Not in CONTEXT_APPROVAL state", str(context.exception))
        
    def test_approve_response_state_error(self):
        """Test approving response in wrong state raises error."""
        with self.assertRaises(Exception) as context:
            self.agent.approve_response()
        self.assertIn("Not in RESPONSE_APPROVAL state", str(context.exception))
        
    async def test_context_approval_flow(self):
        """Test complete context approval flow with state transitions."""
        # Register handlers
        self.agent.add_state_change_handler(self.state_change_handler)
        self.agent.add_response_change_handler(self.response_change_handler)
        
        # Setup LLM response
        def mock_infer(prompt: str, context: str) -> Iterator[str]:
            # Return complete response at once for testing
            yield "<reasoning>test reasoning</reasoning>"
        self.mock_llm.infer.side_effect = mock_infer
        
        # Set initial state
        self.agent._current_state = WebAgentState.CONTEXT_APPROVAL
        
        # Approve context
        self.agent.approve_context()
        
        # Wait for response tasks to complete
        await asyncio.gather(*self.web_server._response_tasks)
        
        # Verify state transitions
        expected_states = [
            WebAgentState.INFERENCE,
            WebAgentState.RESPONSE_APPROVAL
        ]
        self.assertEqual(self.state_changes, expected_states)
        
        # Updated test: Should now expect streaming updates
        self.assertEqual(len(self.response_changes), len(self.received_messages))
        for i, msg in enumerate(self.received_messages):
            self.assertEqual(msg["type"], "RESPONSE_UPDATE")
            self.assertEqual(msg["data"], self.response_changes[i])
        
    def test_response_approval_flow_command(self):
        """Test response approval flow with command execution."""
        # Register handlers
        self.agent.add_state_change_handler(self.state_change_handler)
        
        # Setup command result
        command_result = CommandResult.success()
        mock_command = MockCommand(command_result)
        
        # Mock parser to return command
        self.parser.parse = Mock(return_value=mock_command)
        
        # Set initial state and response
        self.agent._current_state = WebAgentState.RESPONSE_APPROVAL
        self.agent._response = "<stop/>"
        
        # Approve response
        self.agent.approve_response()
        
        # Verify command executed
        self.assertTrue(mock_command.executed)
        
        # Verify state transitions
        expected_states = [
            WebAgentState.UPDATE,
            WebAgentState.CONTEXT_APPROVAL
        ]
        self.assertEqual(self.state_changes, expected_states)
        
        # Verify command result stored
        self.assertEqual(self.agent.command_result, command_result)
        
    def test_response_approval_flow_entry(self):
        """Test response approval flow with entry creation."""
        # Register handlers
        self.agent.add_state_change_handler(self.state_change_handler)
        
        # Set initial state and response
        self.agent._current_state = WebAgentState.RESPONSE_APPROVAL
        self.agent.response = "<reasoning>test reasoning</reasoning>"
        
        # Approve response
        self.agent.approve_response()
        
        # Verify entry added to working memory
        self.assertEqual(self.working_memory.get_entries_count(), 1)
        entry = self.working_memory.get_entries()[0]
        self.assertIsInstance(entry, ReasoningEntry)
        self.assertEqual(entry.content, "test reasoning")
        
        # Verify state transitions
        expected_states = [
            WebAgentState.UPDATE,
            WebAgentState.CONTEXT_APPROVAL
        ]
        self.assertEqual(self.state_changes, expected_states)
        
    async def test_response_validation(self):
        """Test response validation during response setting."""
        # Register handlers
        self.agent.add_response_change_handler(self.response_change_handler)
        
        # Mock validator to fail
        error_message = "Invalid XML"
        self.mock_validator.validate.return_value = error_message
        
        # Set response using async method
        await self.agent._set_response("<invalid>")
        
        # Verify validation error stored
        self.assertEqual(self.agent._validation_error, error_message)
        
        # Verify response updated
        self.assertEqual(self.response_changes, ["<invalid>"])
        
    def tearDown(self):
        """Clean up resources."""
        if hasattr(self, 'agent'):
            del self.agent