"""
Coordinator Runtime: Orchestrates multi-agent workflows with dynamic delegation.

Provides intelligent task routing, specialist discovery, and result synthesis.
"""

from __future__ import annotations

import asyncio
from typing import Any, Dict, List, Optional


class CoordinatorRuntime:
    """
    High-level coordinator that manages task delegation and result synthesis.
    
    Features:
    - Dynamic agent discovery and routing
    - Smart delegation based on task analysis
    - Result aggregation and synthesis
    - Queue depth monitoring for load balancing
    """

    def __init__(
        self,
        message_bus: Any,
        agent_registry: Dict[str, Dict[str, Any]],
        database: Any,
        llm_backend: Any,
        config: Any
    ):
        """
        Initialize coordinator runtime.
        
        Args:
            message_bus: Message bus for task distribution
            agent_registry: Registry of available agents
            database: Database for trace logging
            llm_backend: LLM for planning and synthesis
            config: Environment configuration
        """
        self.bus = message_bus
        self.registry = agent_registry
        self.db = database
        self.llm = llm_backend
        self.config = config

    async def delegate(
        self,
        target_agent: str,
        payload: dict,
        wait: bool = True,
        timeout: int = 60
    ) -> dict:
        """
        Delegate a task to a specialist agent.
        
        Args:
            target_agent: Name of target agent
            payload: Task payload
            wait: If True, wait for response
            timeout: Timeout in seconds (only if wait=True)
        
        Returns:
            Response dict or task_id dict
        """
        task_id = await self.bus.publish_task(target_agent, payload)
        
        if wait:
            response = await self.bus.wait_for_response(task_id, timeout)
            if response:
                return response
            return {
                "task_id": task_id,
                "status": "timeout",
                "error": f"No response after {timeout}s"
            }
        
        return {"task_id": task_id, "status": "queued"}

    async def dispatch(self, task: dict) -> dict:
        """
        Dynamically dispatch a task to the most appropriate agent.
        
        Uses heuristics to infer which specialist should handle the task.
        
        Args:
            task: Task dict with query/description
        
        Returns:
            Response from delegated agent
        """
        available_agents = await self.bus.get_registered_agents()
        target_agent = await self._infer_agent(task, available_agents)
        return await self.delegate(target_agent, task)

    async def dispatch_parallel(
        self,
        subtasks: List[dict],
        timeout: int = 60
    ) -> List[dict]:
        """
        Dispatch multiple subtasks in parallel.
        
        Args:
            subtasks: List of task dicts, each with 'agent' and 'payload'
            timeout: Timeout per task
        
        Returns:
            List of results in same order as subtasks
        """
        coros = [
            self.delegate(
                st.get("agent", "researcher"),
                st.get("payload", {}),
                wait=True,
                timeout=timeout
            )
            for st in subtasks
        ]
        results = await asyncio.gather(*coros, return_exceptions=True)
        
        # Normalize exceptions to error dicts
        normalized = []
        for i, r in enumerate(results):
            if isinstance(r, Exception):
                normalized.append({
                    "status": "error",
                    "error": str(r),
                    "subtask_index": i
                })
            else:
                normalized.append(r)
        
        return normalized

    async def synthesize(
        self,
        task_description: str,
        results: List[dict]
    ) -> str:
        """
        Use LLM to synthesize final response from specialist results.
        
        Args:
            task_description: Original task
            results: List of specialist results
        
        Returns:
            Synthesized response string
        """
        results_summary = "\n\n".join([
            f"Agent: {r.get('agent', 'unknown')}\n"
            f"Status: {r.get('status', 'unknown')}\n"
            f"Result: {r.get('result', r.get('error', 'No result'))}"
            for r in results
        ])

        synthesis_prompt = f"""You are a coordinator synthesizing results from multiple specialist agents.

Original Task: {task_description}

Specialist Results:
{results_summary}

Synthesize these results into a coherent, comprehensive response that addresses the original task.
Focus on:
1. Combining complementary information
2. Resolving any contradictions
3. Providing actionable insights
4. Maintaining clarity and conciseness
"""

        response = await self.llm.generate(
            prompt=synthesis_prompt,
            max_tokens=1500,
            temperature=0.5
        )
        
        return response

    async def _infer_agent(
        self,
        task: dict,
        available_agents: Dict[str, Dict[str, Any]]
    ) -> str:
        """
        Infer which agent should handle this task.
        
        Uses simple heuristics:
        - Keyword matching in task text
        - Agent capabilities from metadata
        - Fallback to 'researcher' if uncertain
        
        Args:
            task: Task dict
            available_agents: Registry of agents with metadata
        
        Returns:
            Agent name
        """
        text = (
            task.get("query") or
            task.get("description") or
            task.get("task") or
            ""
        ).lower()

        # Check if task mentions agent name directly
        for name in available_agents.keys():
            if name in text:
                return name

        # Check for capability keywords
        if any(kw in text for kw in ["research", "search", "find", "investigate"]):
            if "researcher" in available_agents:
                return "researcher"
        
        if any(kw in text for kw in ["analyze", "analysis", "insight"]):
            if "analyst" in available_agents:
                return "analyst"
        
        if any(kw in text for kw in ["summarize", "summary", "condense"]):
            if "summarizer" in available_agents:
                return "summarizer"

        # Fallback: return first available agent or 'researcher'
        if available_agents:
            return list(available_agents.keys())[0]
        
        return "researcher"

    async def get_queue_depths(self) -> Dict[str, int]:
        """
        Get queue depth for all registered agents.
        
        Useful for load balancing and auto-scaling decisions.
        
        Returns:
            Dict mapping agent name to queue depth
        """
        agents = await self.bus.get_registered_agents()
        depths = {}
        
        for agent_name in agents.keys():
            try:
                depth = await self.bus.get_queue_depth(agent_name)
                depths[agent_name] = depth
            except Exception:
                depths[agent_name] = 0
        
        return depths

    def should_scale_up(self, agent_name: str, threshold: int = 10) -> bool:
        """
        Determine if an agent should be scaled up based on queue depth.
        
        Args:
            agent_name: Agent to check
            threshold: Queue depth threshold
        
        Returns:
            True if scaling recommended
        """
        # This would be called by adaptive scaling logic
        # For now, just a placeholder for the scaling feature
        return False
