"""
Agent registration mechanism.
Used to manage multiple Agent instances and supports selecting different Agents by model name.
"""

from typing import Dict, Optional, Type, List, Any
from .BaseAgent import BaseAgent
from config.config import get_config
import threading


class AgentRegistry:
    """Agent registry, managing multiple Agent instances."""
    
    def __init__(self):
        self._agents: Dict[str, BaseAgent] = {}
        self._agent_classes: Dict[str, Type[BaseAgent]] = {}
        self._lock = threading.Lock()
    
    def register_agent_class(self, model_name: str, agent_class: Type[BaseAgent]):
        """
        Register an Agent class.

        Args:
            model_name: Model name, used for the model parameter in the API
            agent_class: Agent class, inheriting from BaseAgent
        """
        with self._lock:
            self._agent_classes[model_name] = agent_class
    
    def _create_agent_instance(
        self, 
        model_name: str,
        name: Optional[str] = None,
        description: Optional[str] = None,
        context_file: Optional[str] = None,
        **kwargs
    ) -> BaseAgent:
        """
        Internal method: create an Agent instance.

        Args:
            model_name: Model name
            name: Agent name
            description: Agent description
            context_file: Context file path
            **kwargs: Other parameters
            
        Returns:
            Agent instance
        """
        if model_name not in self._agent_classes:
            raise ValueError(f"Unknown model: {model_name}")
        
        agent_class = self._agent_classes[model_name]
        
        # Get configuration.
        config = get_config()
        
        # Use default values or passed-in parameters.
        agent_name = name or f"{model_name}-agent"
        agent_description = description or f"Agent instance for {model_name}"
        
        # Create the Agent instance.
        agent = agent_class(
            name=agent_name,
            description=agent_description,
            llm_interface=config.BASIC_INTERFACE,
            context_file=context_file,
            model_name=model_name,  # Pass model_name to the Agent.
            **kwargs
        )
        
        return agent
    
    def get_or_create_agent(
        self, 
        model_name: str,
        **kwargs
    ) -> BaseAgent:
        """
        Get or create an Agent instance (singleton pattern).
        Ensure each model_name corresponds to only one Agent instance.
        
        Args:
            model_name: Model name
            **kwargs: Creation parameters
            
        Returns:
            Agent instance
        """
        with self._lock:
            if model_name not in self._agents:
                self._agents[model_name] = self._create_agent_instance(model_name, **kwargs)
            return self._agents[model_name]
    
    def create_agent(
        self, 
        model_name: str,
        force_new: bool = False,
        **kwargs
    ) -> BaseAgent:
        """
        Create an Agent instance.
        
        Args:
            model_name: Model name
            force_new: Whether to force creation of a new instance, replacing the existing instance
            **kwargs: Other parameters
            
        Returns:
            Agent instance
        """
        with self._lock:
            if force_new or model_name not in self._agents:
                self._agents[model_name] = self._create_agent_instance(model_name, **kwargs)
            return self._agents[model_name]
    
    def get_agent(self, model_name: str) -> Optional[BaseAgent]:
        """
        Get an already-created Agent instance.
        
        Args:
            model_name: Model name
            
        Returns:
            Agent instance or None
        """
        return self._agents.get(model_name)
    
    def list_models(self) -> List[str]:
        """
        List all registered model names.
        
        Returns:
            List of model names
        """
        return list(self._agent_classes.keys())
    
    def list_agents(self) -> List[str]:
        """
        List all created Agent instances.
        
        Returns:
            List of model names for Agent instances
        """
        return list(self._agents.keys())
    
    def clear_agents(self):
        """Clear all Agent instances."""
        with self._lock:
            self._agents.clear()
    
    def remove_agent(self, model_name: str) -> bool:
        """
        Remove an Agent instance.
        
        Args:
            model_name: Model name
            
        Returns:
            Whether removal succeeded
        """
        with self._lock:
            if model_name in self._agents:
                del self._agents[model_name]
                return True
            return False
    
    def get_agent_info(self, model_name: str) -> Optional[Dict[str, Any]]:
        """
        Get Agent information.
        
        Args:
            model_name: Model name
            
        Returns:
            Agent information dictionary or None
        """
        agent = self.get_agent(model_name)
        if agent:
            return {
                "model_name": model_name,
                "name": agent.name,
                "description": agent.description,
                "agent_class": agent.__class__.__name__,
                "toolkit_size": len(agent.toolkit),
                "session_info": agent.get_session_info(),
                "is_singleton": True  # Mark this as a singleton instance.
            }
        return None
    
    def get_all_agents_info(self) -> Dict[str, Dict[str, Any]]:
        """
        Get information for all Agents.
        
        Returns:
            Dictionary of all Agent information
        """
        result = {}
        with self._lock:
            for model_name in self._agents:
                info = self.get_agent_info(model_name)
                if info:
                    result[model_name] = info
        return result
    
    def is_agent_active(self, model_name: str) -> bool:
        """
        Check whether the Agent with the specified model_name has been created.
        
        Args:
            model_name: Model name
            
        Returns:
            Whether it has been created
        """
        return model_name in self._agents
    
    def get_agent_stats(self) -> Dict[str, Any]:
        """
        Get registry statistics.
        
        Returns:
            Statistics dictionary
        """
        with self._lock:
            return {
                "registered_models": len(self._agent_classes),
                "active_agents": len(self._agents),
                "registered_model_list": list(self._agent_classes.keys()),
                "active_agent_list": list(self._agents.keys())
            }


# Global Agent registry instance.
_global_registry = AgentRegistry()


def get_agent_registry() -> AgentRegistry:
    """Get the global Agent registry."""
    return _global_registry


def register_agent(model_name: str, agent_class: Type[BaseAgent]):
    """
    Convenience function for registering an Agent class.
    
    Args:
        model_name: Model name
        agent_class: Agent class
    """
    _global_registry.register_agent_class(model_name, agent_class)


def get_agent(model_name: str, **kwargs) -> BaseAgent:
    """
    Convenience function for getting or creating an Agent instance (singleton pattern).
    
    Args:
        model_name: Model name
        **kwargs: Creation parameters, used only during first creation
        
    Returns:
        Agent instance
    """
    return _global_registry.get_or_create_agent(model_name, **kwargs)


def get_existing_agent(model_name: str) -> Optional[BaseAgent]:
    """
    Convenience function for getting an existing Agent instance.
    
    Args:
        model_name: Model name
        
    Returns:
        Agent instance or None
    """
    return _global_registry.get_agent(model_name)


def create_new_agent(model_name: str, **kwargs) -> BaseAgent:
    """
    Convenience function for forcing creation of a new Agent instance.
    
    Args:
        model_name: Model name
        **kwargs: Creation parameters
        
    Returns:
        Agent instance
    """
    return _global_registry.create_agent(model_name, force_new=True, **kwargs)


def list_available_models() -> List[str]:
    """
    Convenience function for listing all available models.
    
    Returns:
        List of model names
    """
    return _global_registry.list_models()


def get_registry_stats() -> Dict[str, Any]:
    """
    Convenience function for getting registry statistics.
    
    Returns:
        Statistics dictionary
    """
    return _global_registry.get_agent_stats()


def clear_all_agents():
    """
    Convenience function for clearing all Agent instances.
    """
    _global_registry.clear_agents()
    # Also clear BaseAgent's class-level instance cache.
    from .BaseAgent import BaseAgent
    BaseAgent.clear_instances()