"""
Task delegator for multi-agent coordination and delegation.
"""

import asyncio
import logging
from typing import Any, TYPE_CHECKING

from a2a.types import AgentCard, Message

from synqed.agent import Agent
from synqed.client import Client

if TYPE_CHECKING:
    from synqed.orchestrator import Orchestrator

logger = logging.getLogger(__name__)


class TaskDelegator:
    """
    Orchestrates task delegation across multiple agents.
    
    The TaskDelegator analyzes tasks and automatically delegates them to the most
    appropriate agent(s) based on their skills and capabilities.
    
    Example:
        ```python
        delegator = TaskDelegator()
        
        # Register agents
        delegator.register_agent(recipe_agent)
        delegator.register_agent(shopping_agent)
        
        # Submit a task - automatically delegated to the right agent(s)
        result = await delegator.submit_task(
            "Find me a recipe and create a shopping list"
        )
        ```
    """
    
    def __init__(self, orchestrator: "Orchestrator | None" = None):
        """
        Initialize the task delegator.
        
        Args:
            orchestrator: Optional Orchestrator instance for intelligent agent selection.
                         If provided, the delegator will use the orchestrator's LLM-based
                         routing instead of simple skill matching.
        """
        self._agents: dict[str, Agent] = {}
        self._agent_cards: dict[str, AgentCard] = {}
        self._agent_urls: dict[str, str] = {}
        self._clients: dict[str, Client] = {}
        self._orchestrator = orchestrator
    
    def register_agent(
        self,
        agent: Agent | None = None,
        agent_url: str | None = None,
        agent_card: AgentCard | None = None,
    ) -> None:
        """
        Register an agent with the delegator.
        
        Args:
            agent: Agent instance (if running locally)
            agent_url: URL of a remote agent
            agent_card: Pre-loaded agent card
            
        Raises:
            ValueError: If insufficient information is provided
        """
        if agent is not None:
            # Local agent
            agent_id = agent.name
            self._agents[agent_id] = agent
            self._agent_cards[agent_id] = agent.card
            self._agent_urls[agent_id] = agent.url
            logger.info(f"Registered local agent: {agent_id}")
            
            # Also register with orchestrator if available
            if self._orchestrator:
                self._orchestrator.register_agent(
                    agent_card=agent.card,
                    agent_url=agent.url,
                    agent_id=agent_id,
                )
        
        elif agent_url is not None:
            # Remote agent
            agent_id = agent_url
            self._agent_urls[agent_id] = agent_url
            if agent_card:
                self._agent_cards[agent_id] = agent_card
                
                # Register with orchestrator if available
                if self._orchestrator:
                    self._orchestrator.register_agent(
                        agent_card=agent_card,
                        agent_url=agent_url,
                        agent_id=agent_id,
                    )
            logger.info(f"Registered remote agent: {agent_url}")
        
        else:
            raise ValueError(
                "Must provide either an Agent instance, agent_url, or agent_card"
            )
    
    def unregister_agent(self, agent_name_or_url: str) -> None:
        """
        Unregister an agent.
        
        Args:
            agent_name_or_url: Agent name or URL to unregister
        """
        if agent_name_or_url in self._agents:
            del self._agents[agent_name_or_url]
        if agent_name_or_url in self._agent_cards:
            del self._agent_cards[agent_name_or_url]
        if agent_name_or_url in self._agent_urls:
            del self._agent_urls[agent_name_or_url]
        if agent_name_or_url in self._clients:
            del self._clients[agent_name_or_url]
        
        logger.info(f"Unregistered agent: {agent_name_or_url}")
    
    def list_agents(self) -> list[dict[str, Any]]:
        """
        List all registered agents.
        
        Returns:
            List of agent information dictionaries
        """
        agents_info = []
        
        for agent_id in self._agent_urls.keys():
            info = {
                "id": agent_id,
                "url": self._agent_urls.get(agent_id),
                "is_local": agent_id in self._agents,
            }
            
            if agent_id in self._agent_cards:
                card = self._agent_cards[agent_id]
                info["name"] = card.name
                info["description"] = card.description
                info["skills"] = [skill.id for skill in card.skills]
            
            agents_info.append(info)
        
        return agents_info
    
    async def submit_task(
        self,
        task_description: str,
        preferred_agent: str | None = None,
        require_skills: list[str] | None = None,
        use_orchestrator: bool = True,
    ) -> str:
        """
        Submit a task and automatically delegate to appropriate agent(s).
        
        Args:
            task_description: Description of the task to perform
            preferred_agent: Optional agent name/URL to use
            require_skills: Optional list of required skills
            use_orchestrator: Whether to use the orchestrator if available (default: True)
            
        Returns:
            Task result as a string
            
        Raises:
            ValueError: If no suitable agent is found
        """
        # Find the best agent for the task
        agent_id = await self._select_agent(
            task_description,
            preferred_agent,
            require_skills,
            use_orchestrator=use_orchestrator,
        )
        
        if agent_id is None:
            raise ValueError(
                f"No suitable agent found for task. "
                f"Required skills: {require_skills}"
            )
        
        logger.info(f"Delegating task to agent: {agent_id}")
        
        # Get or create client for the agent
        client = await self._get_client(agent_id)
        
        # Send the task
        result = await client.ask(task_description)
        
        return result
    
    async def submit_task_to_multiple(
        self,
        task_description: str,
        agent_ids: list[str] | None = None,
    ) -> dict[str, str]:
        """
        Submit the same task to multiple agents and collect all responses.
        
        Args:
            task_description: Description of the task
            agent_ids: List of agent IDs to send to (None = all agents)
            
        Returns:
            Dictionary mapping agent IDs to their responses
        """
        if agent_ids is None:
            agent_ids = list(self._agent_urls.keys())
        
        if not agent_ids:
            raise ValueError("No agents available")
        
        logger.info(f"Submitting task to {len(agent_ids)} agents")
        
        # Create tasks for all agents
        tasks = []
        for agent_id in agent_ids:
            client = await self._get_client(agent_id)
            tasks.append(client.ask(task_description))
        
        # Wait for all responses
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # Build result dictionary
        response_dict = {}
        for agent_id, result in zip(agent_ids, results):
            if isinstance(result, Exception):
                response_dict[agent_id] = f"Error: {str(result)}"
                logger.error(f"Agent {agent_id} failed: {result}")
            else:
                response_dict[agent_id] = result
        
        return response_dict
    
    async def _select_agent(
        self,
        task_description: str,
        preferred_agent: str | None = None,
        require_skills: list[str] | None = None,
        use_orchestrator: bool = True,
    ) -> str | None:
        """
        Select the best agent for a task.
        
        Args:
            task_description: Task description
            preferred_agent: Preferred agent ID
            require_skills: Required skills
            use_orchestrator: Whether to use orchestrator if available
            
        Returns:
            Selected agent ID or None
        """
        # If preferred agent is specified, use it
        if preferred_agent and preferred_agent in self._agent_urls:
            return preferred_agent
        
        # Use orchestrator if available and requested
        if use_orchestrator and self._orchestrator:
            try:
                logger.info("Using orchestrator for agent selection")
                result = await self._orchestrator.orchestrate(task_description)
                if result.selected_agents:
                    selected = result.selected_agents[0]  # Use the top recommendation
                    logger.info(
                        f"Orchestrator selected: {selected.agent_name} "
                        f"(confidence: {selected.confidence:.2f})"
                    )
                    logger.info(f"Reasoning: {selected.reasoning}")
                    return selected.agent_id
            except Exception as e:
                logger.warning(f"Orchestrator selection failed, falling back to simple matching: {e}")
        
        # Fall back to simple skill-based matching
        # Load all agent cards if needed
        for agent_id in self._agent_urls.keys():
            if agent_id not in self._agent_cards:
                try:
                    client = await self._get_client(agent_id)
                    # Card will be loaded by client
                    if hasattr(client, "_agent_card") and client._agent_card:
                        self._agent_cards[agent_id] = client._agent_card
                except Exception as e:
                    logger.warning(f"Could not load card for {agent_id}: {e}")
        
        # Filter by required skills if specified
        candidates = []
        for agent_id, card in self._agent_cards.items():
            if require_skills:
                agent_skills = {skill.id for skill in card.skills}
                if all(req_skill in agent_skills for req_skill in require_skills):
                    candidates.append(agent_id)
            else:
                candidates.append(agent_id)
        
        # Return first candidate
        if candidates:
            return candidates[0]
        
        # Fall back to any available agent
        if self._agent_urls:
            return next(iter(self._agent_urls.keys()))
        
        return None
    
    async def _get_client(self, agent_id: str) -> Client:
        """Get or create a client for an agent."""
        if agent_id not in self._clients:
            agent_url = self._agent_urls[agent_id]
            agent_card = self._agent_cards.get(agent_id)
            
            self._clients[agent_id] = Client(
                agent_url=agent_url,
                agent_card=agent_card,
            )
        
        return self._clients[agent_id]
    
    async def close_all(self) -> None:
        """Close all client connections."""
        for client in self._clients.values():
            await client.close()
        self._clients.clear()
        logger.info("All clients closed")
    
    def __repr__(self) -> str:
        """String representation."""
        return (
            f"TaskDelegator(registered_agents={len(self._agents)}, "
            f"total_agents={len(self._agent_urls)})"
        )

