import unittest
from datetime import datetime
from unittest.mock import Mock, MagicMock, patch
import xml.etree.ElementTree as ET
import logging

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.base_agent import BaseAgent
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.delete_command import DeleteCommand
from sia.entry import Entry

class TestLogHandler(logging.Handler):
    """Custom logging handler that stores records for test verification."""
    def __init__(self):
        super().__init__()
        self.records = []

    def emit(self, record):
        self.records.append(record)

    def clear(self):
        self.records = []

class MockEntry(Entry):
    """Mock entry class that properly extends Entry."""
    def __init__(self, id: str, timestamp: datetime, update_behavior=None):
        super().__init__(id, timestamp)
        self._update_behavior = update_behavior
        self.update_called = False

    def update(self) -> None:
        self.update_called = True
        if self._update_behavior:
            self._update_behavior()

    def generate_context(self) -> ET.Element:
        elem = ET.Element("mock_entry", {"id": self.id})
        elem.text = "<![CDATA[mock content]]>"
        return elem

class TestBaseAgent(BaseAgent):
    """Concrete implementation of BaseAgent for testing."""
    def run(self) -> None:
        pass

class BaseAgentTest(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()
        
        # 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
        
        # Create test agent with required components
        self.agent = TestBaseAgent(
            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)
        
        # Setup logging
        self.log_handler = TestLogHandler()
        logger = logging.getLogger()
        logger.addHandler(self.log_handler)
        
    def log_handler(self, record):
        """Capture log records for testing."""
        self.log_records.append(record)
    
    def tearDown(self):
        """Clean up resources."""
        # Remove log handler
        logger = logging.getLogger()
        logger.removeHandler(self.log_handler)
        
        if hasattr(self, 'agent'):
            del self.agent
            
    def test_initialization(self):
        """Test agent initialization and component setup."""
        self.assertEqual(self.agent._working_memory, self.working_memory)
        self.assertEqual(self.agent._metrics, self.mock_metrics)
        self.assertEqual(self.agent._llm, self.mock_llm)
        self.assertEqual(self.agent._validator, self.mock_validator)
        self.assertEqual(self.agent._parser, self.parser)
        self.assertEqual(self.agent._action_schema, "test schema")
        
    def test_cleanup(self):
        """Test cleanup on agent deletion."""
        del self.agent
        self.mock_metrics.stop.assert_called_once()
        
    def test_compile_context_empty_memory(self):
        """Test context compilation with empty working memory."""
        context = self.agent._compile_context()
        
        # Parse context and verify structure
        root = ET.fromstring(context)
        self.assertEqual(root.tag, "context")
        
        # Check metrics
        self.assertEqual(root.get("cpu"), "10")
        self.assertEqual(root.get("gpu"), "20")
        
        # Check context size is 0 (empty memory)
        self.assertEqual(root.get("context"), "0")
        
        # Check no memory entries
        self.assertEqual(len(list(root)), 0)
        
    def test_compile_context_with_entries(self):
        # Add test entry to working memory
        entry = ReasoningEntry("test reasoning", "test-id", self.test_timestamp)
        self.agent._working_memory.add_entry(entry)
        
        context = self.agent._compile_context()
        
        # Parse and verify context
        root = ET.fromstring(context)
        
        # Check context size reflects one entry
        self.assertEqual(root.get("context"), "1")
        
        # Verify entry content
        reasoning_elem = root.find("reasoning")
        self.assertIsNotNone(reasoning_elem)
        self.assertEqual(reasoning_elem.get("id"), "test-id")
        # Strip whitespace before comparing
        self.assertEqual(reasoning_elem.text.strip(), "test reasoning")
        
    def test_multiple_entries(self):
        """Test handling multiple entries in working memory."""
        # Add multiple entries
        entries = [
            ReasoningEntry(f"test {i}", f"id-{i}", self.test_timestamp)
            for i in range(3)
        ]
        
        for entry in entries:
            self.agent._working_memory.add_entry(entry)
            
        context = self.agent._compile_context()
        root = ET.fromstring(context)
        
        # Check context size reflects three entries
        self.assertEqual(root.get("context"), "3")
        
        # Verify entry order maintained
        reasoning_elems = root.findall("reasoning")
        self.assertEqual(len(reasoning_elems), 3)
        for i, elem in enumerate(reasoning_elems):
            self.assertEqual(elem.get("id"), f"id-{i}")
            self.assertEqual(elem.text.strip(), f"test {i}")