import asyncio
from typing import Optional, List, Callable
from sia.web_agent import WebAgentState

class MockWebAgent:
    """Mock WebAgent for testing."""
    def __init__(self):
        self._current_state: Optional[WebAgentState] = None
        self.context: Optional[str] = None
        self.response: Optional[str] = None
        self._validation_error: Optional[str] = None
        self._state_handlers: List[Callable] = []
        self._response_handlers: List[Callable] = []
        
    @property
    def current_state(self):
        return self._current_state
        
    def add_state_change_handler(self, handler):
        self._state_handlers.append(handler)
        
    def add_response_change_handler(self, handler):
        self._response_handlers.append(handler)
        
    async def _notify_state_handlers(self, new_state: WebAgentState):
        """Notify all state handlers of a state change."""
        for handler in self._state_handlers:
            await handler(new_state)
            
    async def _notify_response_handlers(self, response: str):
        """Notify all response handlers of a response change."""
        for handler in self._response_handlers:
            await handler(response)
            
    async def _set_response(self, response: str):
        """Set response and notify handlers."""
        self.response = response
        await self._notify_response_handlers(response)
            
    async def approve_context(self):
        """Handle context approval with immediate state updates."""
        # Change to INFERENCE state and notify
        self._current_state = WebAgentState.INFERENCE
        for handler in self._state_handlers:
            await handler(WebAgentState.INFERENCE)
        
        # Set response and notify
        self.response = "<reasoning>test</reasoning>"
        for handler in self._response_handlers:
            await handler(self.response)
        
        # Change to RESPONSE_APPROVAL state and notify
        self._current_state = WebAgentState.RESPONSE_APPROVAL
        for handler in self._state_handlers:
            await handler(WebAgentState.RESPONSE_APPROVAL)