"""Registry for managing multiple LangGraph agents."""
from typing import Dict, Callable, Any, Optional
from langchain_core.language_models.chat_models import BaseChatModel
from graphs.workflows.orchestrated import create_orchestrated_graph
from graphs.workflows.orchestrated_v2 import create_orchestrated_graph_v2
from graphs.workflows.agent import create_agent_graph

# Canonical registry key when the client omits `agent` (API metadata uses .upper()).
_DEFAULT_AGENT_KEY = "agent"


class AgentRegistry:
    """
    Registry for managing multiple agent graph builders.
    
    This allows for easy addition of new agent ids without modifying
    the API layer. Each agent id maps to a graph builder function.
    """
    
    def __init__(self):
        """Initialize the agent registry with default agents."""
        self._agent_builders: Dict[str, Callable[[BaseChatModel], Any]] = {
            "v1": create_orchestrated_graph,      # V1 (unchanged)
            "v2": create_orchestrated_graph_v2,   # V2 (isolated)
            "agent": create_agent_graph,                # Agent (isolated)
        }
        self._descriptions: Dict[str, str] = {
            "v1": "Current production orchestrated workflow",
            "v2": "Isolated V2 workflow",
            "agent": "ReAct agent workflow (default)",
        }
    
    def register_agent(
        self,
        agent: str,
        builder: Callable[[BaseChatModel], Any],
        description: str = ""
    ) -> None:
        """
        Register a new agent builder.
        
        Args:
            agent: Agent identifier (string)
            builder: Function that takes an LLM and returns a compiled graph
            description: Description of the agent
        """
        agent_key = agent.strip().lower()
        self._agent_builders[agent_key] = builder
        if description:
            self._descriptions[agent_key] = description

    def get_builder_for_request(
        self,
        agent: Optional[str],
    ) -> Callable[[BaseChatModel], Any]:
        """Resolve graph builder from request.

        If agent is missing, defaults to the ReAct agent (`agent`).
        """
        agent_key = (agent or _DEFAULT_AGENT_KEY).strip().lower()
        if agent_key not in self._agent_builders:
            raise ValueError(
                f"Agent '{agent}' not implemented. "
                f"Available agents: {list(self._agent_builders.keys())}"
            )
        return self._agent_builders[agent_key]

    def resolve_agent_id(self, agent: Optional[str]) -> str:
        """Return canonical agent id used in API metadata."""
        return (agent or _DEFAULT_AGENT_KEY).strip().upper()
    
    def list_agents(self) -> list[dict]:
        """
        List all registered agents with their information.
        
        Returns:
            List of agent information dictionaries
        """
        return [
            {
                "type": key.upper(),
                "name": key.upper(),
                "description": self._descriptions.get(key, "No description available"),
                "available": True,
            }
            for key in sorted(self._agent_builders.keys())
        ]


# Singleton instance
agent_registry = AgentRegistry()