from datetime import datetime, timezone
from enum import Enum, auto
from sys import exit
from typing import Callable, Dict, List

from .base_agent import BaseAgent
from .command import Command
from .iteration_logger import IterationLogger
from .llm_engine import LlmEngine
from .response_buffer import ResponseBuffer
from .response_parser import ResponseParser
from .system_metrics import SystemMetrics
from .working_memory import WorkingMemory

class AgentState(Enum):
    IDLE = auto()
    INFERENCE = auto()
    PROCESSING_RESPONSE = auto()

class WebAgent(BaseAgent):
    def __init__(
            self,
            system_prompt: str,
            action_schema: str, 
            working_memory: WorkingMemory,
            metrics: SystemMetrics,
            llms: Dict[str, LlmEngine],
            parser: ResponseParser,
            iteration_logger: IterationLogger,
        ):
        super().__init__(
            system_prompt,
            action_schema,
            working_memory,
            metrics,
            parser
        )

        self._llms = llms
        self._selected_llm = list(llms.keys())[0]
        self._state: AgentState = AgentState.IDLE

        self._iteration_logger = iteration_logger
        self._response_buffer: ResponseBuffer = ResponseBuffer()
        
        self._state_change_handlers: List[Callable[[AgentState], None]] = []
        self._selected_llm_change_handlers: List[Callable[[str], None]] = []
        self._context_change_handlers: List[Callable[[Dict], None]] = []

        self._update_compiled_context()
        self._working_memory.add_change_handler(self._update_compiled_context)
        self._parser._io_buffer.add_change_handler(lambda *_: self._update_compiled_context())
    
    @property
    def response_buffer(self) -> ResponseBuffer:
        return self._response_buffer

    @property
    def llms(self) -> List[str]:
        return self._llms.keys()
    
    @property
    def active_llm(self) -> str:
        return self._selected_llm
    
    @active_llm.setter
    def active_llm(self, llm_name: str) -> None:
        if llm_name not in self._llms:
            raise ValueError(f"Invalid LLM name: {llm_name}")
        if self._selected_llm == llm_name:
            return
        self._selected_llm = llm_name
        self._update_compiled_context()
        self._notify_selected_llm_change()
    
    def add_selected_llm_change_handler(self, handler: Callable[[str], None]) -> None:
        self._selected_llm_change_handlers.append(handler)
    
    @property
    def state(self) -> AgentState:
        return self._state

    def add_state_change_handler(self, handler: Callable[[AgentState], None]) -> None:
        self._state_change_handlers.append(handler)

    @property
    def context_info(self) -> Dict:
        return dict(self._compiled_context.attrib)

    def add_context_change_handler(self, handler: Callable[[Dict], None]) -> None:
        self._context_change_handlers.append(handler)

    def run_inference(self) -> None:
        """Start inference on specified LLM"""
        if self._state != AgentState.IDLE:
            raise RuntimeError(f"Not ready for inference")
        self._update_state(AgentState.INFERENCE)
        llm = self._llms[self.active_llm]
        response_token_iter = llm.infer(
            self.system_prompt,
            self._compiled_context,
            self._response_buffer.get_text()
        )
        for token in response_token_iter:
            self._response_buffer.append_text(token)
        self._update_state(AgentState.IDLE)

    def stop_inference(self) -> None:
        """Stop ongoing inference for specified LLM"""
        self._llms[self.active_llm].restart()
        self._update_state(AgentState.IDLE)

    def approve_response(self) -> None:
        """Process approved response from specified LLM"""
        self._update_state(AgentState.PROCESSING_RESPONSE)
        timestamp = datetime.now(timezone.utc)
        self._iteration_logger.log_iteration(
            timestamp,
            self._compiled_context,
            self._response_buffer.get_text()
        )
        parse_result = self._parser.parse(timestamp, self._response_buffer.get_text())
        self._response_buffer.clear()
        if isinstance(parse_result, Command):
            result = parse_result.execute(self._working_memory)
            self._command_result = result
            if result.should_stop:
                exit(42)
            else:
                self._working_memory.update()
        else:
            parse_result.update()
            self._working_memory.update()
            self._working_memory.add_entry(parse_result)
        self._update_state(AgentState.IDLE)

    def _update_state(self, state: AgentState) -> None:
        """Update LLM state and notify handlers"""
        self._state = state
        for handler in self._state_change_handlers:
            handler(state)
    
    def _update_compiled_context(self):
        self._compiled_context = self._compile_context(self._llms[self.active_llm])
        self._notify_context_change()

    def _notify_context_change(self):
        """Notify all context change handlers."""
        for handler in self._context_change_handlers:
            handler(self.context_info)

    def _notify_selected_llm_change(self):
        """Notify all selected LLM change handlers."""
        for handler in self._selected_llm_change_handlers:
            handler(self._selected_llm)