"""
Multi-Agent Orchestration Framework for Toxo.

This module implements an advanced multi-agent system with:
- Agent lifecycle management
- Task distribution and scheduling
- Inter-agent communication
- Agent coordination patterns (hierarchical, peer-to-peer, consensus)
- Dynamic agent scaling
- Fault tolerance and recovery
- Performance monitoring and optimization
"""

import asyncio
import json
import uuid
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any, Union, Callable, Tuple, Set
from enum import Enum
from datetime import datetime, timedelta
import logging
from pathlib import Path
import threading
from contextlib import asynccontextmanager

import numpy as np
from pydantic import BaseModel, Field

from ..core.config import ToxoConfig
from ..core.load_balancer import LoadBalancer
from ..core.caching import CacheManager
from typing import Any
from ..utils.logger import get_logger
from ..utils.exceptions import ToxoError
from ..utils.performance import PerformanceTracker


class AgentType(Enum):
    """Types of agents in the system."""
    COORDINATOR = "coordinator"
    WORKER = "worker"
    SPECIALIST = "specialist"
    MONITOR = "monitor"
    PLANNER = "planner"
    EXECUTOR = "executor"
    VALIDATOR = "validator"


class AgentStatus(Enum):
    """Agent lifecycle status."""
    INITIALIZING = "initializing"
    ACTIVE = "active"
    BUSY = "busy"
    IDLE = "idle"
    ERROR = "error"
    STOPPING = "stopping"
    STOPPED = "stopped"


class TaskStatus(Enum):
    """Task execution status."""
    PENDING = "pending"
    ASSIGNED = "assigned"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class CoordinationPattern(Enum):
    """Agent coordination patterns."""
    HIERARCHICAL = "hierarchical"
    PEER_TO_PEER = "peer_to_peer"
    CONSENSUS = "consensus"
    PIPELINE = "pipeline"
    BROADCAST = "broadcast"
    LEADER_FOLLOWER = "leader_follower"


class MessageType(Enum):
    """Inter-agent message types."""
    TASK_ASSIGNMENT = "task_assignment"
    TASK_RESULT = "task_result"
    STATUS_UPDATE = "status_update"
    COORDINATION = "coordination"
    HEARTBEAT = "heartbeat"
    ERROR_REPORT = "error_report"
    RESOURCE_REQUEST = "resource_request"


@dataclass
class AgentCapability:
    """Defines an agent's capabilities."""
    name: str
    description: str
    input_types: List[str]
    output_types: List[str]
    max_concurrent_tasks: int = 1
    estimated_duration: Optional[float] = None  # seconds
    required_resources: Dict[str, Any] = field(default_factory=dict)


@dataclass
class AgentMetrics:
    """Performance metrics for an agent."""
    tasks_completed: int = 0
    tasks_failed: int = 0
    total_execution_time: float = 0.0
    average_execution_time: float = 0.0
    cpu_usage: float = 0.0
    memory_usage: float = 0.0
    success_rate: float = 1.0
    last_activity: Optional[datetime] = None
    
    def update_completion(self, execution_time: float, success: bool):
        """Update metrics after task completion."""
        if success:
            self.tasks_completed += 1
            self.total_execution_time += execution_time
            self.average_execution_time = self.total_execution_time / self.tasks_completed
        else:
            self.tasks_failed += 1
        
        total_tasks = self.tasks_completed + self.tasks_failed
        self.success_rate = self.tasks_completed / total_tasks if total_tasks > 0 else 1.0
        self.last_activity = datetime.now()


@dataclass
class Message:
    """Inter-agent message."""
    id: str
    sender_id: str
    receiver_id: str
    message_type: MessageType
    content: Dict[str, Any]
    timestamp: datetime = field(default_factory=datetime.now)
    priority: int = 0  # Higher numbers = higher priority
    requires_response: bool = False
    correlation_id: Optional[str] = None


@dataclass
class Task:
    """Task to be executed by agents."""
    id: str
    type: str
    description: str
    input_data: Dict[str, Any]
    output_data: Optional[Dict[str, Any]] = None
    status: TaskStatus = TaskStatus.PENDING
    assigned_agent_id: Optional[str] = None
    created_at: datetime = field(default_factory=datetime.now)
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    priority: int = 0
    max_retries: int = 3
    retry_count: int = 0
    dependencies: List[str] = field(default_factory=list)  # Task IDs
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    @property
    def execution_time(self) -> Optional[float]:
        """Get task execution time in seconds."""
        if self.started_at and self.completed_at:
            return (self.completed_at - self.started_at).total_seconds()
        return None


class Agent(ABC):
    """Abstract base class for all agents."""
    
    def __init__(
        self,
        agent_id: str,
        agent_type: AgentType,
        capabilities: List[AgentCapability],
        config: Optional[Dict[str, Any]] = None
    ):
        self.agent_id = agent_id
        self.agent_type = agent_type
        self.capabilities = {cap.name: cap for cap in capabilities}
        self.config = config or {}
        self.status = AgentStatus.INITIALIZING
        self.metrics = AgentMetrics()
        self.logger = get_logger(f"{__name__}.{agent_id}")
        
        # Communication
        self.message_queue: asyncio.Queue = asyncio.Queue()
        self.orchestrator: Optional['AgentOrchestrator'] = None
        
        # Task management
        self.current_tasks: Dict[str, Task] = {}
        self.max_concurrent_tasks = max(cap.max_concurrent_tasks for cap in capabilities)
        
        # Monitoring
        self.last_heartbeat = datetime.now()
        self._running = False
        self._task_handlers = {}
    
    async def initialize(self):
        """Initialize the agent."""
        try:
            await self._setup()
            self.status = AgentStatus.ACTIVE
            self.logger.info(f"Agent {self.agent_id} initialized successfully")
        except Exception as e:
            self.status = AgentStatus.ERROR
            self.logger.error(f"Failed to initialize agent {self.agent_id}: {e}")
            raise
    
    async def start(self):
        """Start the agent's main loop."""
        self._running = True
        self.status = AgentStatus.ACTIVE
        
        # Start main processing loop
        await asyncio.gather(
            self._message_loop(),
            self._heartbeat_loop(),
            self._task_monitor_loop()
        )
    
    async def stop(self):
        """Stop the agent gracefully."""
        self.status = AgentStatus.STOPPING
        self._running = False
        
        # Wait for current tasks to complete
        while self.current_tasks:
            await asyncio.sleep(0.1)
        
        await self._cleanup()
        self.status = AgentStatus.STOPPED
        self.logger.info(f"Agent {self.agent_id} stopped")
    
    async def send_message(self, message: Message):
        """Send message to another agent."""
        if self.orchestrator:
            await self.orchestrator.route_message(message)
        else:
            self.logger.warning(f"No orchestrator available to send message: {message.id}")
    
    async def receive_message(self, message: Message):
        """Receive message from another agent."""
        await self.message_queue.put(message)
    
    async def assign_task(self, task: Task) -> bool:
        """Assign a task to this agent."""
        if len(self.current_tasks) >= self.max_concurrent_tasks:
            return False
        
        if not self.can_handle_task(task):
            return False
        
        task.assigned_agent_id = self.agent_id
        task.status = TaskStatus.ASSIGNED
        self.current_tasks[task.id] = task
        
        # Start task execution
        asyncio.create_task(self._execute_task(task))
        return True
    
    def can_handle_task(self, task: Task) -> bool:
        """Check if agent can handle the given task."""
        return task.type in self.capabilities
    
    @abstractmethod
    async def _setup(self):
        """Setup agent-specific resources."""
        pass
    
    @abstractmethod
    async def _cleanup(self):
        """Cleanup agent-specific resources."""
        pass
    
    @abstractmethod
    async def execute_task(self, task: Task) -> Dict[str, Any]:
        """Execute a specific task (implemented by subclasses)."""
        pass
    
    async def _message_loop(self):
        """Main message processing loop."""
        while self._running:
            try:
                # Wait for message with timeout
                message = await asyncio.wait_for(
                    self.message_queue.get(),
                    timeout=1.0
                )
                await self._handle_message(message)
            except asyncio.TimeoutError:
                continue
            except Exception as e:
                self.logger.error(f"Error in message loop: {e}")
    
    async def _heartbeat_loop(self):
        """Send periodic heartbeats."""
        while self._running:
            try:
                if self.orchestrator:
                    heartbeat = Message(
                        id=str(uuid.uuid4()),
                        sender_id=self.agent_id,
                        receiver_id="orchestrator",
                        message_type=MessageType.HEARTBEAT,
                        content={
                            "status": self.status.value,
                            "active_tasks": len(self.current_tasks),
                            "metrics": {
                                "tasks_completed": self.metrics.tasks_completed,
                                "tasks_failed": self.metrics.tasks_failed,
                                "success_rate": self.metrics.success_rate
                            }
                        }
                    )
                    await self.send_message(heartbeat)
                
                self.last_heartbeat = datetime.now()
                await asyncio.sleep(30)  # Heartbeat every 30 seconds
                
            except Exception as e:
                self.logger.error(f"Error in heartbeat loop: {e}")
    
    async def _task_monitor_loop(self):
        """Monitor task execution and handle timeouts."""
        while self._running:
            try:
                current_time = datetime.now()
                
                for task_id, task in list(self.current_tasks.items()):
                    # Check for task timeout
                    if task.started_at:
                        elapsed = (current_time - task.started_at).total_seconds()
                        capability = self.capabilities.get(task.type)
                        
                        if capability and capability.estimated_duration:
                            timeout = capability.estimated_duration * 3  # 3x estimated time
                            if elapsed > timeout:
                                self.logger.warning(f"Task {task_id} timed out")
                                await self._handle_task_timeout(task)
                
                await asyncio.sleep(10)  # Check every 10 seconds
                
            except Exception as e:
                self.logger.error(f"Error in task monitor loop: {e}")
    
    async def _handle_message(self, message: Message):
        """Handle incoming message."""
        try:
            if message.message_type == MessageType.TASK_ASSIGNMENT:
                # Handle task assignment
                task_data = message.content.get('task')
                if task_data:
                    task = Task(**task_data)
                    success = await self.assign_task(task)
                    
                    if message.requires_response:
                        response = Message(
                            id=str(uuid.uuid4()),
                            sender_id=self.agent_id,
                            receiver_id=message.sender_id,
                            message_type=MessageType.STATUS_UPDATE,
                            content={"task_assigned": success, "task_id": task.id},
                            correlation_id=message.id
                        )
                        await self.send_message(response)
            
            elif message.message_type == MessageType.STATUS_UPDATE:
                # Handle status updates from other agents
                await self._handle_status_update(message)
            
            elif message.message_type == MessageType.COORDINATION:
                # Handle coordination messages
                await self._handle_coordination_message(message)
            
            # Add more message type handlers as needed
            
        except Exception as e:
            self.logger.error(f"Error handling message {message.id}: {e}")
    
    async def _execute_task(self, task: Task):
        """Execute a task and handle the lifecycle."""
        start_time = datetime.now()
        task.started_at = start_time
        task.status = TaskStatus.IN_PROGRESS
        
        try:
            self.logger.info(f"Starting task {task.id}: {task.description}")
            
            # Execute the task
            result = await self.execute_task(task)
            
            # Update task
            task.output_data = result
            task.status = TaskStatus.COMPLETED
            task.completed_at = datetime.now()
            
            # Update metrics
            execution_time = (task.completed_at - start_time).total_seconds()
            self.metrics.update_completion(execution_time, True)
            
            self.logger.info(f"Completed task {task.id} in {execution_time:.2f}s")
            
            # Notify orchestrator of completion
            if self.orchestrator:
                completion_message = Message(
                    id=str(uuid.uuid4()),
                    sender_id=self.agent_id,
                    receiver_id="orchestrator",
                    message_type=MessageType.TASK_RESULT,
                    content={
                        "task_id": task.id,
                        "status": task.status.value,
                        "result": result,
                        "execution_time": execution_time
                    }
                )
                await self.send_message(completion_message)
        
        except Exception as e:
            self.logger.error(f"Task {task.id} failed: {e}")
            
            task.status = TaskStatus.FAILED
            task.completed_at = datetime.now()
            
            # Update metrics
            execution_time = (task.completed_at - start_time).total_seconds()
            self.metrics.update_completion(execution_time, False)
            
            # Handle retry logic
            if task.retry_count < task.max_retries:
                task.retry_count += 1
                task.status = TaskStatus.PENDING
                task.started_at = None
                self.logger.info(f"Retrying task {task.id} (attempt {task.retry_count + 1})")
                
                # Reschedule task
                await asyncio.sleep(2 ** task.retry_count)  # Exponential backoff
                asyncio.create_task(self._execute_task(task))
                return
            
            # Notify orchestrator of failure
            if self.orchestrator:
                failure_message = Message(
                    id=str(uuid.uuid4()),
                    sender_id=self.agent_id,
                    receiver_id="orchestrator",
                    message_type=MessageType.ERROR_REPORT,
                    content={
                        "task_id": task.id,
                        "error": str(e),
                        "retry_count": task.retry_count
                    }
                )
                await self.send_message(failure_message)
        
        finally:
            # Remove task from current tasks
            self.current_tasks.pop(task.id, None)
            
            # Update agent status
            if not self.current_tasks:
                self.status = AgentStatus.IDLE
            elif len(self.current_tasks) < self.max_concurrent_tasks:
                self.status = AgentStatus.ACTIVE
            else:
                self.status = AgentStatus.BUSY
    
    async def _handle_task_timeout(self, task: Task):
        """Handle task timeout."""
        self.logger.warning(f"Task {task.id} timed out after {task.execution_time}s")
        
        task.status = TaskStatus.FAILED
        task.completed_at = datetime.now()
        
        # Remove from current tasks
        self.current_tasks.pop(task.id, None)
        
        # Update metrics
        if task.execution_time:
            self.metrics.update_completion(task.execution_time, False)
    
    async def _handle_status_update(self, message: Message):
        """Handle status update from another agent."""
        # Default implementation - can be overridden by subclasses
        pass
    
    async def _handle_coordination_message(self, message: Message):
        """Handle coordination message from another agent."""
        # Default implementation - can be overridden by subclasses
        pass


class LLMAgent(Agent):
    """Agent specialized for LLM-based tasks."""
    
    def __init__(
        self,
        agent_id: str,
        llm_client: Any,
        capabilities: Optional[List[AgentCapability]] = None,
        config: Optional[Dict[str, Any]] = None
    ):
        if capabilities is None:
            capabilities = [
                AgentCapability(
                    name="text_generation",
                    description="Generate text based on prompts",
                    input_types=["prompt"],
                    output_types=["text"],
                    max_concurrent_tasks=3,
                    estimated_duration=5.0
                ),
                AgentCapability(
                    name="text_analysis",
                    description="Analyze and extract insights from text",
                    input_types=["text"],
                    output_types=["analysis"],
                    max_concurrent_tasks=5,
                    estimated_duration=3.0
                )
            ]
        
        super().__init__(agent_id, AgentType.SPECIALIST, capabilities, config)
        # Must support `await llm_client.generate(prompt=..., temperature=..., max_tokens=...)`
        self.llm_client = llm_client
    
    async def _setup(self):
        """Setup LLM agent resources."""
        # Test LLM connection
        await self.llm_client.generate("Hello", temperature=0.1)
    
    async def _cleanup(self):
        """Cleanup LLM agent resources."""
        pass
    
    async def execute_task(self, task: Task) -> Dict[str, Any]:
        """Execute LLM-based task."""
        if task.type == "text_generation":
            prompt = task.input_data.get("prompt", "")
            temperature = task.input_data.get("temperature", 0.7)
            max_tokens = task.input_data.get("max_tokens")
            
            response = await self.llm_client.generate(
                prompt=prompt,
                temperature=temperature,
                max_tokens=max_tokens
            )
            
            return {"generated_text": response}
        
        elif task.type == "text_analysis":
            text = task.input_data.get("text", "")
            analysis_type = task.input_data.get("analysis_type", "general")
            
            if analysis_type == "sentiment":
                prompt = f"Analyze the sentiment of the following text and return a JSON with sentiment (positive/negative/neutral) and confidence score:\n\n{text}"
            elif analysis_type == "summary":
                prompt = f"Provide a concise summary of the following text:\n\n{text}"
            elif analysis_type == "keywords":
                prompt = f"Extract key topics and keywords from the following text as a JSON list:\n\n{text}"
            else:
                prompt = f"Analyze the following text and provide insights:\n\n{text}"
            
            response = await self.llm_client.generate(prompt, temperature=0.3)
            
            return {"analysis": response, "analysis_type": analysis_type}
        
        else:
            raise ToxoError(f"Unknown task type: {task.type}")
    
    def get_state(self) -> Dict[str, Any]:
        """Get the current state for serialization."""
        return {
            "agent_id": self.agent_id,
            "llm_client_available": getattr(self, 'llm_client', None) is not None,
            "capabilities_count": len(getattr(self, 'capabilities', {})),
            "config_available": getattr(self, 'config', None) is not None,
            "status": getattr(self, 'status', AgentStatus.INITIALIZING).value,
            "system_type": "LLMAgent"
        }


class TaskScheduler:
    """Advanced task scheduler for the orchestrator."""
    
    def __init__(self):
        self.pending_tasks: asyncio.PriorityQueue = asyncio.PriorityQueue()
        self.completed_tasks: Dict[str, Task] = {}
        self.task_dependencies: Dict[str, Set[str]] = {}  # task_id -> dependent task_ids
        self.logger = get_logger(f"{__name__}.scheduler")
    
    async def submit_task(self, task: Task):
        """Submit a task for scheduling."""
        # Calculate task priority (higher number = higher priority)
        priority = -task.priority  # Negative for min-heap behavior
        
        await self.pending_tasks.put((priority, task.created_at, task))
        self.logger.info(f"Task {task.id} submitted for scheduling")
        
        # Track dependencies
        if task.dependencies:
            for dep_task_id in task.dependencies:
                if dep_task_id not in self.task_dependencies:
                    self.task_dependencies[dep_task_id] = set()
                self.task_dependencies[dep_task_id].add(task.id)
    
    async def get_next_task(self) -> Optional[Task]:
        """Get the next task to be executed."""
        try:
            while not self.pending_tasks.empty():
                priority, created_at, task = await self.pending_tasks.get()
                
                # Check if dependencies are satisfied
                if self._dependencies_satisfied(task):
                    return task
                else:
                    # Put task back in queue
                    await self.pending_tasks.put((priority, created_at, task))
                    await asyncio.sleep(0.1)  # Avoid busy waiting
            
            return None
            
        except asyncio.QueueEmpty:
            return None
    
    def _dependencies_satisfied(self, task: Task) -> bool:
        """Check if all task dependencies are satisfied."""
        for dep_task_id in task.dependencies:
            dep_task = self.completed_tasks.get(dep_task_id)
            if not dep_task or dep_task.status != TaskStatus.COMPLETED:
                return False
        return True
    
    def mark_task_completed(self, task: Task):
        """Mark a task as completed and update dependent tasks."""
        self.completed_tasks[task.id] = task
        
        # Check if any dependent tasks can now be scheduled
        dependent_tasks = self.task_dependencies.get(task.id, set())
        for dep_task_id in dependent_tasks:
            # Find and potentially reschedule dependent tasks
            pass  # Implementation depends on how you want to handle this
    
    def get_pending_count(self) -> int:
        """Get number of pending tasks."""
        return self.pending_tasks.qsize()


class AgentOrchestrator:
    """Main orchestrator for managing multiple agents."""
    
    def __init__(
        self,
        config: ToxoConfig,
        coordination_pattern: CoordinationPattern = CoordinationPattern.HIERARCHICAL
    ):
        self.config = config
        self.coordination_pattern = coordination_pattern
        self.logger = get_logger(f"{__name__}.orchestrator")
        
        # Agent management
        self.agents: Dict[str, Agent] = {}
        self.agent_capabilities: Dict[str, List[str]] = {}  # agent_id -> capability names
        
        # Task management
        self.scheduler = TaskScheduler()
        self.active_tasks: Dict[str, Task] = {}
        
        # Communication
        self.message_router = MessageRouter()
        
        # Monitoring
        self.performance_tracker = PerformanceTracker()
        self.load_balancer = LoadBalancer()
        
        # State
        self._running = False
        self._coordinator_task = None
    
    async def start(self):
        """Start the orchestrator."""
        self._running = True
        self.logger.info("Starting agent orchestrator")
        
        # Start main coordination loop
        self._coordinator_task = asyncio.create_task(self._coordination_loop())
        
        # Start all agents
        agent_tasks = [agent.start() for agent in self.agents.values()]
        if agent_tasks:
            await asyncio.gather(*agent_tasks, return_exceptions=True)
    
    async def stop(self):
        """Stop the orchestrator gracefully."""
        self._running = False
        self.logger.info("Stopping agent orchestrator")
        
        # Stop all agents
        stop_tasks = [agent.stop() for agent in self.agents.values()]
        if stop_tasks:
            await asyncio.gather(*stop_tasks, return_exceptions=True)
        
        # Cancel coordinator task
        if self._coordinator_task:
            self._coordinator_task.cancel()
            try:
                await self._coordinator_task
            except asyncio.CancelledError:
                pass
    
    async def register_agent(self, agent: Agent):
        """Register a new agent with the orchestrator."""
        self.agents[agent.agent_id] = agent
        self.agent_capabilities[agent.agent_id] = list(agent.capabilities.keys())
        agent.orchestrator = self
        
        await agent.initialize()
        self.logger.info(f"Registered agent {agent.agent_id} with capabilities: {list(agent.capabilities.keys())}")
    
    async def unregister_agent(self, agent_id: str):
        """Unregister an agent."""
        if agent_id in self.agents:
            agent = self.agents[agent_id]
            await agent.stop()
            del self.agents[agent_id]
            del self.agent_capabilities[agent_id]
            self.logger.info(f"Unregistered agent {agent_id}")
    
    async def submit_task(self, task: Task) -> str:
        """Submit a task for execution."""
        await self.scheduler.submit_task(task)
        self.active_tasks[task.id] = task
        self.logger.info(f"Submitted task {task.id}: {task.description}")
        return task.id
    
    async def get_task_status(self, task_id: str) -> Optional[TaskStatus]:
        """Get the status of a task."""
        task = self.active_tasks.get(task_id)
        return task.status if task else None
    
    async def get_task_result(self, task_id: str) -> Optional[Dict[str, Any]]:
        """Get the result of a completed task."""
        task = self.active_tasks.get(task_id)
        if task and task.status == TaskStatus.COMPLETED:
            return task.output_data
        return None
    
    async def route_message(self, message: Message):
        """Route message between agents."""
        await self.message_router.route_message(message, self.agents)
    
    async def _coordination_loop(self):
        """Main coordination loop."""
        while self._running:
            try:
                # Get next task to schedule
                task = await self.scheduler.get_next_task()
                if task:
                    await self._assign_task(task)
                
                # Handle agent health checks
                await self._check_agent_health()
                
                # Balance load if needed
                await self._balance_load()
                
                await asyncio.sleep(1)  # Coordination interval
                
            except Exception as e:
                self.logger.error(f"Error in coordination loop: {e}")
                await asyncio.sleep(5)  # Longer sleep on error
    
    async def _assign_task(self, task: Task) -> bool:
        """Assign a task to the best available agent."""
        # Find capable agents
        capable_agents = [
            agent for agent in self.agents.values()
            if agent.can_handle_task(task) and agent.status in [AgentStatus.ACTIVE, AgentStatus.IDLE]
        ]
        
        if not capable_agents:
            # Put task back in queue
            await self.scheduler.submit_task(task)
            return False
        
        # Select best agent based on current load and performance
        best_agent = self._select_best_agent(capable_agents, task)
        
        # Assign task
        success = await best_agent.assign_task(task)
        if success:
            self.logger.info(f"Assigned task {task.id} to agent {best_agent.agent_id}")
        else:
            # Put task back in queue
            await self.scheduler.submit_task(task)
        
        return success
    
    def _select_best_agent(self, agents: List[Agent], task: Task) -> Agent:
        """Select the best agent for a task based on various criteria."""
        if len(agents) == 1:
            return agents[0]
        
        # Score agents based on multiple factors
        agent_scores = []
        
        for agent in agents:
            score = 0.0
            
            # Factor 1: Current load (lower is better)
            load_factor = 1.0 - (len(agent.current_tasks) / agent.max_concurrent_tasks)
            score += load_factor * 0.4
            
            # Factor 2: Success rate (higher is better)
            score += agent.metrics.success_rate * 0.3
            
            # Factor 3: Average execution time (lower is better, normalized)
            if agent.metrics.average_execution_time > 0:
                # Normalize by capability estimated duration
                capability = agent.capabilities.get(task.type)
                if capability and capability.estimated_duration:
                    time_factor = capability.estimated_duration / agent.metrics.average_execution_time
                    score += min(time_factor, 1.0) * 0.2
            else:
                score += 0.2  # Give benefit of doubt for new agents
            
            # Factor 4: Idle time preference
            if agent.status == AgentStatus.IDLE:
                score += 0.1
            
            agent_scores.append((agent, score))
        
        # Return agent with highest score
        best_agent = max(agent_scores, key=lambda x: x[1])[0]
        return best_agent
    
    async def _check_agent_health(self):
        """Check health of all agents."""
        current_time = datetime.now()
        
        for agent in list(self.agents.values()):
            # Check if agent is responsive
            time_since_heartbeat = (current_time - agent.last_heartbeat).total_seconds()
            
            if time_since_heartbeat > 120:  # 2 minutes without heartbeat
                self.logger.warning(f"Agent {agent.agent_id} appears unresponsive")
                agent.status = AgentStatus.ERROR
                
                # Could implement automatic recovery here
                # await self._recover_agent(agent)
    
    async def _balance_load(self):
        """Balance load across agents if needed."""
        if self.coordination_pattern == CoordinationPattern.HIERARCHICAL:
            # In hierarchical mode, coordinator handles load balancing
            
            # Get agent utilization
            agent_utilization = {}
            for agent_id, agent in self.agents.items():
                utilization = len(agent.current_tasks) / agent.max_concurrent_tasks
                agent_utilization[agent_id] = utilization
            
            # Identify overloaded and underloaded agents
            overloaded = [aid for aid, util in agent_utilization.items() if util > 0.8]
            underloaded = [aid for aid, util in agent_utilization.items() if util < 0.3]
            
            # Could implement task migration logic here
            if overloaded and underloaded:
                self.logger.info(f"Load imbalance detected: {len(overloaded)} overloaded, {len(underloaded)} underloaded")
    
    def get_system_metrics(self) -> Dict[str, Any]:
        """Get comprehensive system metrics."""
        total_agents = len(self.agents)
        active_agents = len([a for a in self.agents.values() if a.status == AgentStatus.ACTIVE])
        total_tasks = len(self.active_tasks)
        pending_tasks = self.scheduler.get_pending_count()
        
        # Aggregate agent metrics
        total_completed = sum(a.metrics.tasks_completed for a in self.agents.values())
        total_failed = sum(a.metrics.tasks_failed for a in self.agents.values())
        avg_success_rate = np.mean([a.metrics.success_rate for a in self.agents.values()]) if self.agents else 0
        
        return {
            "orchestrator": {
                "coordination_pattern": self.coordination_pattern.value,
                "running": self._running
            },
            "agents": {
                "total": total_agents,
                "active": active_agents,
                "by_status": {
                    status.value: len([a for a in self.agents.values() if a.status == status])
                    for status in AgentStatus
                }
            },
            "tasks": {
                "total_active": total_tasks,
                "pending": pending_tasks,
                "completed": total_completed,
                "failed": total_failed,
                "success_rate": avg_success_rate
            },
            "performance": self.performance_tracker.get_summary() if hasattr(self, 'performance_tracker') else {}
        }
    
    def get_state(self) -> Dict[str, Any]:
        """Get the current state for serialization."""
        return {
            "config_available": self.config is not None,
            "coordination_pattern": self.coordination_pattern.value if hasattr(self, 'coordination_pattern') else "unknown",
            "agents_count": len(getattr(self, 'agents', {})),
            "active_tasks_count": len(getattr(self, 'active_tasks', {})),
            "system_type": "AgentOrchestrator"
        }


class MessageRouter:
    """Routes messages between agents."""
    
    def __init__(self):
        self.logger = get_logger(f"{__name__}.message_router")
        self.message_history: List[Message] = []
        self.max_history = 1000
    
    async def route_message(self, message: Message, agents: Dict[str, Agent]):
        """Route a message to its destination."""
        try:
            # Store message in history
            self.message_history.append(message)
            if len(self.message_history) > self.max_history:
                self.message_history.pop(0)
            
            # Route to specific agent or broadcast
            if message.receiver_id == "all" or message.receiver_id == "broadcast":
                # Broadcast to all agents except sender
                for agent_id, agent in agents.items():
                    if agent_id != message.sender_id:
                        await agent.receive_message(message)
            else:
                # Route to specific agent
                target_agent = agents.get(message.receiver_id)
                if target_agent:
                    await target_agent.receive_message(message)
                else:
                    self.logger.warning(f"Unknown recipient for message {message.id}: {message.receiver_id}")
        
        except Exception as e:
            self.logger.error(f"Failed to route message {message.id}: {e}")


# Factory functions
def create_llm_agent(agent_id: str, llm_client: Any, **config) -> LLMAgent:
    """Create an LLM agent."""
    return LLMAgent(agent_id, llm_client, config=config)


def create_orchestrator(config: ToxoConfig, **kwargs) -> AgentOrchestrator:
    """Create an agent orchestrator."""
    return AgentOrchestrator(config, **kwargs)


async def create_multi_agent_system(
    config: ToxoConfig,
    llm_client: Any,
    num_llm_agents: int = 3
) -> AgentOrchestrator:
    """Create a complete multi-agent system."""
    orchestrator = create_orchestrator(config)
    
    # Create LLM agents
    for i in range(num_llm_agents):
        agent = create_llm_agent(f"llm_agent_{i}", llm_client)
        await orchestrator.register_agent(agent)
    
    return orchestrator 