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

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.entry.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."""
        self.mock_llm = Mock(spec=LlmEngine)
        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 system metrics
        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
        }
        
        # Mock LLM response generator
        def mock_infer(prompt: str, context: str):
            yield "<reasoning>test reasoning</reasoning>"
            
        self.mock_llm.infer.side_effect = mock_infer
        self.mock_llm.token_count.return_value = 100
        self.mock_llm.token_limit.return_value = 1000
        
        # Create agent with all 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
        )
        
        # 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, validation_error: 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.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._llm_change_handlers), 0)
        self.assertEqual(len(self.agent._response_change_handlers), 0)
        
    def test_handler_registration(self):
        """Test adding state and response change handlers."""
        self.agent.add_llm_change_handler(self.state_change_handler)
        self.agent.add_response_change_handler(self.response_change_handler)
        
        self.assertEqual(len(self.agent._llm_change_handlers), 1)
        self.assertEqual(len(self.agent._response_change_handlers), 1)
        
        # Adding same handler twice should not duplicate
        self.agent.add_llm_change_handler(self.state_change_handler)
        self.agent.add_response_change_handler(self.response_change_handler)
        
        self.assertEqual(len(self.agent._llm_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."""
        self.agent._state = WebAgentState.UPDATE
        
        with self.assertRaises(Exception) as context:
            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))
        
    def test_context_approval_flow(self):
        """Test complete context approval flow with state transitions."""
        self.agent.add_llm_change_handler(self.state_change_handler)
        self.agent.add_response_change_handler(self.response_change_handler)
        
        self.agent.approve_context()
        
        expected_states = [
            WebAgentState.INFERENCE,
            WebAgentState.RESPONSE_APPROVAL
        ]
        self.assertEqual(self.state_changes, expected_states)
        self.assertEqual(self.response_changes[-1], "<reasoning>test reasoning</reasoning>")
        
    def test_response_approval_flow_command(self):
        """Test response approval flow with command execution."""
        self.agent.add_llm_change_handler(self.state_change_handler)
        
        command_result = CommandResult.success()
        mock_command = MockCommand(command_result)
        self.parser.parse = Mock(return_value=mock_command)
        
        self.agent._state = WebAgentState.RESPONSE_APPROVAL
        self.agent._response = "<stop/>"
        
        self.agent.approve_response()
        
        self.assertTrue(mock_command.executed)
        
        expected_states = [
            WebAgentState.UPDATE,
            WebAgentState.CONTEXT_APPROVAL
        ]
        self.assertEqual(self.state_changes, expected_states)
        self.assertEqual(self.agent.command_result, command_result)
        
    def test_response_approval_flow_entry(self):
        """Test response approval flow with entry creation."""
        self.agent.add_llm_change_handler(self.state_change_handler)
        
        self.agent._state = WebAgentState.RESPONSE_APPROVAL
        self.agent._set_response("<reasoning>test reasoning</reasoning>")
        
        self.agent.approve_response()
        time.sleep(1)
        
        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")
        
        expected_states = [
            WebAgentState.UPDATE,
            WebAgentState.CONTEXT_APPROVAL
        ]
        self.assertEqual(self.state_changes, expected_states)
        
    def test_response_validation(self):
        """Test response validation during response setting."""
        self.agent.add_response_change_handler(self.response_change_handler)
        error_message = "Invalid XML"
        self.mock_validator.validate.return_value = error_message
        self.agent._set_response("<invalid>")
        self.assertEqual(self.agent.validation_error, error_message)
        self.assertEqual(self.response_changes, ["<invalid>"])