from collections import defaultdict
from datetime import datetime, timezone
from enum import Enum, auto
from sys import exit
from threading import Lock
from typing import Callable, Dict, List, Optional

from .base_agent import BaseAgent
from .command import Command
from .command_result import CommandResult
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
from .xml_validator import XMLValidator

class LlmState(Enum):
    IDLE = auto()
    INFERENCE = auto()

class WebAgent(BaseAgent):
    def __init__(
            self,
            system_prompt: str,
            action_schema: str, 
            working_memory: WorkingMemory,
            metrics: SystemMetrics,
            llms: Dict[str, LlmEngine],
            validator: XMLValidator,
            parser: ResponseParser,
            iteration_logger: IterationLogger,
        ):
        super().__init__(
            system_prompt,
            action_schema,
            working_memory,
            metrics,
            validator,
            parser
        )
        self._llms = llms
        self._iteration_logger = iteration_logger
        self._llm_states: Dict[str, LlmState] = {name: LlmState.IDLE for name in llms}
        self._response_buffer: ResponseBuffer = ResponseBuffer()
        self._validation_error: Optional[str] = None
        self._command_result: Optional[CommandResult] = None
        self._context = self._compile_context(next(iter(self._llms.values())))
        self._stop_flags: Dict[str, bool] = {name: False for name in llms}
        
        # Locks
        self._llm_lock = Lock()
        self._output_lock = Lock()
        
        # Event handlers
        self._llm_change_handlers: List[Callable[[str, LlmState], None]] = []
        self._token_handlers: List[Callable[[str, str], None]] = []
        self._context_change_handlers: List[Callable[[str, bool], None]] = []

        # Change handlers
        self._working_memory.add_change_handler(self._handle_memory_update)

    @property
    def llms(self) -> Dict[str, LlmState]:
        """Get current state of all LLMs"""
        with self._llm_lock:
            return self._llm_states.copy()
    
    @property
    def response_buffer(self) -> ResponseBuffer:
        return self._response_buffer

    @property
    def context(self) -> str:
        return self._context

    @property
    def command_result(self) -> Optional[CommandResult]:
        return self._command_result

    @property
    def validation_error(self) -> Optional[str]:
        return self._validation_error

    def add_llm_change_handler(self, handler: Callable[[str, LlmState], None]) -> None:
        """Add handler for LLM state changes"""
        if handler not in self._llm_change_handlers:
            self._llm_change_handlers.append(handler)

    def add_context_change_handler(self, handler: Callable[[str, bool], None]) -> None:
        """Add handler for context changes"""
        if handler not in self._context_change_handlers:
            self._context_change_handlers.append(handler)

    def modify_context(self, context: str, generated: bool = False) -> None:
        """Update context and reset all LLM states"""
        with self._llm_lock:
            self._context = context
            self._response_buffer.clear()
            for llm_name in self._llms:
                self._set_llm_state(llm_name, LlmState.IDLE)
            
            for handler in self._context_change_handlers:
                handler(context, generated)

    def run_inference(self, llm_name: str) -> None:
        """Start inference on specified LLM"""
        if llm_name not in self._llms:
            raise ValueError(f"Unknown LLM: {llm_name}")
        with self._llm_lock:
            if self._llm_states[llm_name] != LlmState.IDLE:
                raise RuntimeError(f"LLM {llm_name} is not ready for inference")
            self._set_llm_state(llm_name, LlmState.INFERENCE)
            self._stop_flags[llm_name] = False
        llm = self._llms[llm_name]
        def should_stop() -> bool:
            return self._stop_flags[llm_name]
        response_token_iter = llm.infer(self.system_prompt, self.context, self._response_buffer.get_text(), should_stop)
        for token in response_token_iter:
            print(token, end='', flush=True)
            with self._output_lock:
                self._response_buffer.append_text(token)
        with self._llm_lock:
            self._set_llm_state(llm_name, LlmState.IDLE)

    def stop_inference(self, llm_name: str) -> None:
        """Stop ongoing inference for specified LLM"""
        if llm_name not in self._llms:
            raise ValueError(f"Unknown LLM: {llm_name}")
        self._stop_flags[llm_name] = True

    def approve_response(self) -> None:
        """Process approved response from specified LLM"""
        if self.llms.get(llm_name) != LlmState.IDLE:
            return
        timestamp = datetime.now(timezone.utc)
        self._iteration_logger.log_iteration(timestamp, self._context, self._response_buffer.get_text())
        parse_result = self._parser.parse(timestamp, self._response_buffer.get_text())
        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)

    def _set_llm_state(self, llm_name: str, state: LlmState) -> None:
        """Update LLM state and notify handlers"""
        self._llm_states[llm_name] = state
        for handler in self._llm_change_handlers:
            handler(llm_name, state)

    def _handle_memory_update(self) -> None:
        """Handle memory updates and update context"""
        context = self._compile_context(next(iter(self._llms.values())))
        self.modify_context(context, True)