from dataclasses import dataclass
from threading import Thread, Event
from typing import Callable, TypeAlias
from .web_agent import WebAgent, AgentState

@dataclass
class AutoApproverConfig:
    context_enabled: bool
    response_enabled: bool
    context_timeout: float
    response_timeout: float

ConfigChangeHandler: TypeAlias = Callable[[AutoApproverConfig], None]

class AutoApprover:
    """
    Handles automatic approvals for WebAgent states after configurable timeouts.
    """
    
    def __init__(self, agent: WebAgent):
        """
        Initialize auto approver with a WebAgent instance.
        """
        self.agent = agent
        self._config = AutoApproverConfig(
            context_enabled=False,
            response_enabled=False,
            context_timeout=1.0,
            response_timeout=1.0,
        )
        self._stop_event = Event()
        self._context_thread: Thread | None = None
        self._response_thread: Thread | None = None
        self._config_change_handlers: list[ConfigChangeHandler] = []
        
        self._previous_state = agent.state
        self.agent.add_state_change_handler(self._handle_state_change)

    @property
    def config(self) -> AutoApproverConfig:
        return self._config

    @config.setter
    def config(self, config: AutoApproverConfig) -> None:
        if config == self._config:
            return
        if config.context_timeout != self._config.context_timeout:
            if config.context_enabled and self._config.context_enabled:
                raise ValueError("Cannot change context timeout while enabled")
        if config.response_timeout != self._config.response_timeout:
            if config.response_enabled and self._config.response_enabled:
                raise ValueError("Cannot change response timeout while enabled")
        self._stop_context_thread()
        self._stop_response_thread()
        # If both context and response are enabled, start response approval only.
        # Better to run inference with a full prefix than approve an emmpty response.
        if config.response_enabled and not self._config.response_enabled:
            # response approval was just enabled
            self._config = config
            self._start_response_thread()
        elif config.context_enabled and not self._config.context_enabled:
            # context approval was just enabled
            self._config = config
            self._start_context_thread()
        self._config = config
        self._notify_config_change()

    def add_config_change_handler(self, handler: ConfigChangeHandler) -> None:
        self._config_change_handlers.append(handler)

    def _notify_config_change(self) -> None:
        current_config = self.config
        for handler in self._config_change_handlers:
            handler(current_config)

    def _handle_state_change(self, state: AgentState) -> None:
        if state == AgentState.IDLE:
            if self._previous_state == AgentState.INFERENCE:
                # finished inference
                if self._config.response_enabled:
                    self._start_response_thread()
            elif self._previous_state == AgentState.PROCESSING_RESPONSE:
                # finished processing response
                if self._config.context_enabled:
                    self._start_context_thread()
        self._previous_state = state

    def _stop_context_thread(self) -> None:
        if self._context_thread:
            self._stop_event.set()
            self._context_thread = None
            self._stop_event.clear()

    def _stop_response_thread(self) -> None:
        if self._response_thread:
            self._stop_event.set()
            self._response_thread = None
            self._stop_event.clear()

    def _start_context_thread(self) -> None:
        self._context_thread = Thread(target=self._context_approval_thread)
        self._context_thread.start()

    def _start_response_thread(self) -> None:
        self._response_thread = Thread(target=self._response_approval_thread)
        self._response_thread.start()

    def _context_approval_thread(self) -> None:
        if self._stop_event.wait(self._config.context_timeout):
            return
        if self._config.context_enabled:
            self.agent.run_inference()

    def _response_approval_thread(self) -> None:
        if self._stop_event.wait(self._config.response_timeout):
            return
        if self._config.response_enabled:
            self.agent.approve_response()