"""
Agent registry and decorators for AgentChain.
"""

from typing import Dict, Type, Optional, Callable, Any, Union
import logging
import functools

from .base import Agent, Orchestrator, AgentContext, AgentResult


logger = logging.getLogger(__name__)


class AgentRegistry:
    """
    Global registry of available agents.
    
    Agents can be registered via decorator or direct registration.
    """
    
    _instance: Optional["AgentRegistry"] = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._agents: Dict[str, Type[Agent]] = {}
            cls._instance._orchestrators: Dict[str, Type[Orchestrator]] = {}
        return cls._instance
    
    @classmethod
    def get_instance(cls) -> "AgentRegistry":
        """Get the singleton instance."""
        return cls()
    
    def register(
        self,
        agent_class: Type[Agent],
        name: Optional[str] = None
    ) -> Type[Agent]:
        """Register an agent class."""
        agent_name = name or agent_class.name
        
        if issubclass(agent_class, Orchestrator):
            self._orchestrators[agent_name] = agent_class
            logger.debug(f"Registered orchestrator: {agent_name}")
        else:
            self._agents[agent_name] = agent_class
            logger.debug(f"Registered agent: {agent_name}")
        
        return agent_class
    
    def get(self, name: str) -> Optional[Type[Agent]]:
        """Get an agent class by name."""
        return self._agents.get(name) or self._orchestrators.get(name)
    
    def get_orchestrator(self, name: str) -> Optional[Type[Orchestrator]]:
        """Get an orchestrator class by name."""
        return self._orchestrators.get(name)
    
    def list_agents(self) -> Dict[str, Type[Agent]]:
        """List all registered agents."""
        return dict(self._agents)
    
    def list_orchestrators(self) -> Dict[str, Type[Orchestrator]]:
        """List all registered orchestrators."""
        return dict(self._orchestrators)
    
    def clear(self) -> None:
        """Clear all registered agents."""
        self._agents.clear()
        self._orchestrators.clear()


# Global registry instance
_registry = AgentRegistry.get_instance()


def agent(
    name: Optional[str] = None,
    role: str = "Agent",
    model: Optional[str] = None,
    system_prompt: str = "",
    is_orchestrator: bool = False
) -> Callable:
    """
    Decorator to register an agent.
    
    Can be used on a class or function.
    
    Class usage:
        @agent(name="my_agent", role="Does stuff")
        class MyAgent(Agent):
            def execute(self, context):
                ...
    
    Function usage:
        @agent(name="simple_agent")
        def simple_agent(context: AgentContext) -> AgentResult:
            return AgentResult.ok("done")
    """
    def decorator(obj: Union[Type[Agent], Callable]) -> Union[Type[Agent], Callable]:
        if isinstance(obj, type) and issubclass(obj, Agent):
            # Class decorator
            if name:
                obj.name = name
            if role:
                obj.role = role
            if system_prompt:
                obj.system_prompt = system_prompt
            
            _registry.register(obj, name)
            return obj
        else:
            # Function decorator - wrap in an Agent class
            func = obj
            agent_name = name or func.__name__
            
            @functools.wraps(func)
            class FunctionAgent(Orchestrator if is_orchestrator else Agent):
                name = agent_name
                role = role
                system_prompt = system_prompt
                
                def __init__(self, **kwargs):
                    super().__init__(**kwargs)
                    if model:
                        self.model = model
                
                def execute(self, context: AgentContext) -> AgentResult:
                    result = func(context, self)
                    if isinstance(result, AgentResult):
                        return result
                    return AgentResult.ok(result)
            
            FunctionAgent.__name__ = f"{agent_name.title()}Agent"
            _registry.register(FunctionAgent, agent_name)
            return FunctionAgent
    
    return decorator


def orchestrator(
    name: Optional[str] = None,
    role: str = "Orchestrator",
    model: Optional[str] = None,
    system_prompt: str = ""
) -> Callable:
    """
    Decorator to register an orchestrator agent.
    
    Same as @agent but with is_orchestrator=True.
    """
    return agent(
        name=name,
        role=role,
        model=model,
        system_prompt=system_prompt,
        is_orchestrator=True
    )


def get_registry() -> AgentRegistry:
    """Get the global agent registry."""
    return _registry


def get_agent(name: str) -> Optional[Type[Agent]]:
    """Get an agent class by name from the global registry."""
    return _registry.get(name)


def list_agents() -> Dict[str, Type[Agent]]:
    """List all registered agents."""
    return _registry.list_agents()
