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

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_parser import ResponseParser
from .system_metrics import SystemMetrics
from .working_memory import WorkingMemory
from .xml_validator import XMLValidator

class LlmState(Enum):
    NO_OUTPUT = auto()
    INFERENCE = auto()
    OUTPUT = 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.NO_OUTPUT for name in llms}
        self._llm_outputs: Dict[str, str] = defaultdict(str)
        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]] = []

        # Working memory change handler
        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 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_token_handler(self, handler: Callable[[str, str], None]) -> None:
        """Add handler for new tokens"""
        if handler not in self._token_handlers:
            self._token_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._llm_outputs.clear()
            for llm_name in self._llms:
                self._set_llm_state(llm_name, LlmState.NO_OUTPUT)
            
            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.NO_OUTPUT:
                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, should_stop)
        
        with self._output_lock:
            self._llm_outputs[llm_name] = ""

        for token in response_token_iter:
            with self._output_lock:
                self._llm_outputs[llm_name] += token
            
            for handler in self._token_handlers:
                handler(llm_name, token)

        with self._llm_lock:
            self._set_llm_state(llm_name, LlmState.OUTPUT)

    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 get_output(self, llm_name: str) -> str:
        """Get complete output for specified LLM"""
        if llm_name not in self._llms:
            raise ValueError(f"Unknown LLM: {llm_name}")
            
        with self._output_lock:
            return self._llm_outputs[llm_name]

    def approve_response(self, llm_name: str, response: str) -> None:
        """Process approved response from specified LLM"""
        if llm_name not in self._llms:
            raise ValueError(f"Unknown LLM: {llm_name}")

        timestamp = datetime.now()
        self._iteration_logger.log_iteration(timestamp, self._context, response)
        parse_result = self._parser.parse(timestamp, response)
        if isinstance(parse_result, Command):
            result = parse_result.execute(self._working_memory)
            self._command_result = result
            if not result.should_stop:
                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)