"""
Orchestrator - Intelligent agent routing using LLMs.

The Orchestrator analyzes tasks and uses an LLM to intelligently decide which
agents should handle them based on their skills and capabilities.
"""

import json
import logging
from dataclasses import dataclass
from enum import Enum
from typing import Any

from a2a.types import AgentCard

logger = logging.getLogger(__name__)


class LLMProvider(str, Enum):
    """Supported LLM providers."""
    OPENAI = "openai"
    ANTHROPIC = "anthropic"
    GOOGLE = "google"


@dataclass
class AgentSelection:
    """Result of agent selection by the orchestrator."""
    agent_id: str
    agent_name: str
    confidence: float  # 0.0 to 1.0
    reasoning: str
    recommended_skills: list[str]


@dataclass
class OrchestrationResult:
    """Complete result from the orchestrator."""
    task: str
    selected_agents: list[AgentSelection]
    execution_plan: str
    alternative_agents: list[AgentSelection] | None = None


class Orchestrator:
    """
    Intelligent task orchestration using LLMs.
    
    The Orchestrator uses an LLM to analyze tasks and intelligently route them
    to the most appropriate agent(s) based on their capabilities.
    
    Example:
        ```python
        # Initialize with OpenAI
        orchestrator = Orchestrator(
            provider=LLMProvider.OPENAI,
            api_key="sk-...",
            model="gpt-4o"
        )
        
        # Register agents
        orchestrator.register_agent(recipe_agent)
        orchestrator.register_agent(shopping_agent)
        
        # Get routing decision
        result = await orchestrator.orchestrate(
            "Find me a recipe for pasta and create a shopping list"
        )
        
        print(f"Selected agents: {[a.agent_name for a in result.selected_agents]}")
        print(f"Execution plan: {result.execution_plan}")
        ```
    """
    
    def __init__(
        self,
        provider: LLMProvider,
        api_key: str,
        model: str,
        base_url: str | None = None,
        temperature: float = 0.7,
        max_tokens: int = 2000,
    ):
        """
        Initialize the orchestrator with LLM configuration.
        
        Args:
            provider: LLM provider (openai, anthropic, google)
            api_key: API key for the LLM provider
            model: Model name (e.g., "gpt-4o", "claude-3-5-sonnet-20241022", "gemini-pro")
            base_url: Optional base URL for custom endpoints (reserved for future use)
            temperature: LLM temperature (0.0 to 1.0)
            max_tokens: Maximum tokens for LLM response
            
        Raises:
            ValueError: If provider is not supported or missing dependencies
        """
        self.provider = provider
        self.api_key = api_key
        self.model = model
        self.base_url = base_url
        self.temperature = temperature
        self.max_tokens = max_tokens
        
        # Storage for registered agents
        self._agents: dict[str, tuple[AgentCard, str]] = {}  # agent_id -> (card, url)
        
        # Initialize LLM client
        self._llm_client = self._init_llm_client()
    
    def _init_llm_client(self) -> Any:
        """Initialize the LLM client based on the provider."""
        if self.provider == LLMProvider.OPENAI:
            try:
                from openai import AsyncOpenAI
                return AsyncOpenAI(api_key=self.api_key)
            except ImportError:
                raise ValueError(
                    "OpenAI provider requires 'openai' package. "
                    "Install it with: pip install openai"
                )
        
        elif self.provider == LLMProvider.ANTHROPIC:
            try:
                from anthropic import AsyncAnthropic
                return AsyncAnthropic(api_key=self.api_key)
            except ImportError:
                raise ValueError(
                    "Anthropic provider requires 'anthropic' package. "
                    "Install it with: pip install anthropic"
                )
        
        elif self.provider == LLMProvider.GOOGLE:
            try:
                import google.generativeai as genai
                genai.configure(api_key=self.api_key)
                return genai
            except ImportError:
                raise ValueError(
                    "Google provider requires 'google-generativeai' package. "
                    "Install it with: pip install google-generativeai"
                )
        
        else:
            raise ValueError(f"Unsupported LLM provider: {self.provider}")
    
    def register_agent(
        self,
        agent_card: AgentCard,
        agent_url: str,
        agent_id: str | None = None,
    ) -> None:
        """
        Register an agent with the orchestrator.
        
        Args:
            agent_card: Agent's card containing skills and capabilities
            agent_url: URL where the agent can be reached
            agent_id: Optional custom ID (defaults to agent name)
        """
        agent_id = agent_id or agent_card.name
        self._agents[agent_id] = (agent_card, agent_url)
        logger.info(f"Registered agent: {agent_id}")
    
    def unregister_agent(self, agent_id: str) -> None:
        """
        Unregister an agent.
        
        Args:
            agent_id: ID of the agent to unregister
        """
        if agent_id in self._agents:
            del self._agents[agent_id]
            logger.info(f"Unregistered agent: {agent_id}")
    
    def list_agents(self) -> list[dict[str, Any]]:
        """
        List all registered agents.
        
        Returns:
            List of agent information dictionaries
        """
        agents_info = []
        for agent_id, (card, url) in self._agents.items():
            agents_info.append({
                "id": agent_id,
                "name": card.name,
                "description": card.description,
                "skills": [{"id": skill.id, "name": skill.name, "description": skill.description} 
                          for skill in card.skills],
                "url": url,
            })
        return agents_info
    
    async def orchestrate(
        self,
        task: str,
        context: dict[str, Any] | None = None,
        max_agents: int = 3,
    ) -> OrchestrationResult:
        """
        Analyze a task and determine which agent(s) should handle it.
        
        Args:
            task: The task description from the user
            context: Optional additional context about the task
            max_agents: Maximum number of agents to select
            
        Returns:
            OrchestrationResult with selected agents and execution plan
            
        Raises:
            ValueError: If no agents are registered or LLM call fails
        """
        if not self._agents:
            raise ValueError("No agents registered with the orchestrator")
        
        logger.info(f"Orchestrating task: {task[:100]}...")
        
        # Build the prompt for the LLM
        prompt = self._build_orchestration_prompt(task, context, max_agents)
        
        # Call the LLM
        response = await self._call_llm(prompt)
        
        # Parse the response
        result = self._parse_orchestration_response(task, response)
        
        logger.info(f"Selected {len(result.selected_agents)} agent(s) for task")
        return result
    
    def _build_orchestration_prompt(
        self,
        task: str,
        context: dict[str, Any] | None,
        max_agents: int,
    ) -> str:
        """Build the prompt for the LLM."""
        # Gather all agent information
        agents_info = []
        for agent_id, (card, url) in self._agents.items():
            agent_desc = {
                "id": agent_id,
                "name": card.name,
                "description": card.description,
                "skills": [
                    {
                        "id": skill.id,
                        "name": skill.name,
                        "description": skill.description,
                        "tags": skill.tags,
                        "examples": skill.examples or [],
                    }
                    for skill in card.skills
                ],
            }
            agents_info.append(agent_desc)
        
        # Build the prompt
        prompt = f"""You are an intelligent task orchestrator. Your job is to analyze a user's task and determine which agent(s) should handle it based on their skills and capabilities.

USER TASK:
{task}

ADDITIONAL CONTEXT:
{json.dumps(context or {}, indent=2)}

AVAILABLE AGENTS:
{json.dumps(agents_info, indent=2)}

INSTRUCTIONS:
1. Analyze the user's task carefully
2. Review each agent's skills and capabilities
3. Select the most appropriate agent(s) to handle this task (maximum: {max_agents})
4. For each selected agent, explain WHY it was chosen and which skills are relevant
5. Provide an execution plan describing how the agent(s) should work together
6. Rate your confidence in each selection (0.0 to 1.0)

RESPONSE FORMAT (must be valid JSON):
{{
  "selected_agents": [
    {{
      "agent_id": "agent_identifier",
      "agent_name": "Agent Name",
      "confidence": 0.95,
      "reasoning": "Explain why this agent was selected",
      "recommended_skills": ["skill_id_1", "skill_id_2"]
    }}
  ],
  "execution_plan": "Describe how the selected agent(s) should work together to complete the task",
  "alternative_agents": [
    {{
      "agent_id": "alternative_agent_id",
      "agent_name": "Alternative Agent Name", 
      "confidence": 0.6,
      "reasoning": "Why this agent could also work",
      "recommended_skills": ["skill_id"]
    }}
  ]
}}

Return ONLY the JSON response, no additional text.
"""
        return prompt
    
    async def _call_llm(self, prompt: str) -> str:
        """Call the configured LLM with the prompt."""
        try:
            if self.provider == LLMProvider.OPENAI:
                response = await self._llm_client.chat.completions.create(
                    model=self.model,
                    messages=[
                        {"role": "system", "content": "You are a task orchestration expert."},
                        {"role": "user", "content": prompt}
                    ],
                    temperature=self.temperature,
                    max_tokens=self.max_tokens,
                )
                return response.choices[0].message.content
            
            elif self.provider == LLMProvider.ANTHROPIC:
                response = await self._llm_client.messages.create(
                    model=self.model,
                    max_tokens=self.max_tokens,
                    temperature=self.temperature,
                    messages=[
                        {"role": "user", "content": prompt}
                    ]
                )
                return response.content[0].text
            
            elif self.provider == LLMProvider.GOOGLE:
                model = self._llm_client.GenerativeModel(self.model)
                response = await model.generate_content_async(
                    prompt,
                    generation_config={
                        "temperature": self.temperature,
                        "max_output_tokens": self.max_tokens,
                    }
                )
                return response.text
            
            else:
                raise ValueError(f"Provider {self.provider} not implemented")
        
        except Exception as e:
            logger.error(f"LLM call failed: {e}")
            raise ValueError(f"Failed to get orchestration response from LLM: {e}") from e
    
    def _parse_orchestration_response(
        self,
        task: str,
        response: str,
    ) -> OrchestrationResult:
        """Parse the LLM's JSON response into an OrchestrationResult."""
        try:
            # Extract JSON from response (in case LLM adds extra text)
            response = response.strip()
            if response.startswith("```json"):
                response = response[7:]
            if response.startswith("```"):
                response = response[3:]
            if response.endswith("```"):
                response = response[:-3]
            response = response.strip()
            
            data = json.loads(response)
            
            # Parse selected agents
            selected_agents = [
                AgentSelection(
                    agent_id=agent["agent_id"],
                    agent_name=agent["agent_name"],
                    confidence=float(agent["confidence"]),
                    reasoning=agent["reasoning"],
                    recommended_skills=agent.get("recommended_skills", []),
                )
                for agent in data.get("selected_agents", [])
            ]
            
            # Parse alternative agents (optional)
            alternative_agents = None
            if "alternative_agents" in data and data["alternative_agents"]:
                alternative_agents = [
                    AgentSelection(
                        agent_id=agent["agent_id"],
                        agent_name=agent["agent_name"],
                        confidence=float(agent["confidence"]),
                        reasoning=agent["reasoning"],
                        recommended_skills=agent.get("recommended_skills", []),
                    )
                    for agent in data["alternative_agents"]
                ]
            
            return OrchestrationResult(
                task=task,
                selected_agents=selected_agents,
                execution_plan=data.get("execution_plan", "No execution plan provided"),
                alternative_agents=alternative_agents,
            )
        
        except (json.JSONDecodeError, KeyError, ValueError) as e:
            logger.error(f"Failed to parse orchestration response: {e}")
            logger.debug(f"Raw response: {response}")
            raise ValueError(
                f"Failed to parse orchestration response. "
                f"Response may not be valid JSON: {e}"
            ) from e
    
    def __repr__(self) -> str:
        """String representation."""
        return (
            f"Orchestrator(provider={self.provider.value}, "
            f"model={self.model}, "
            f"registered_agents={len(self._agents)})"
        )

