"""
Coordinator Agent - Master Orchestrator for Toxo Multi-Agent System
Implements task decomposition, resource allocation, and agent coordination.
"""

import asyncio
import logging
from typing import Dict, List, Any, Optional, Set, Tuple, Union
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum
import uuid
import json

from ..utils.logger import get_logger
from ..utils.exceptions import CoordinationError, AgentError
from .communication import AgentMessage, MessageBus, MessageType, Priority
from .base import BaseAgent, AgentStatus, AgentCapability


class TaskType(Enum):
    """Types of tasks that can be coordinated."""
    TRAINING = "training"
    DOCUMENT_PROCESSING = "document_processing"
    KNOWLEDGE_EXTRACTION = "knowledge_extraction"
    QUALITY_ASSESSMENT = "quality_assessment"
    USER_INTERACTION = "user_interaction"
    PERFORMANCE_OPTIMIZATION = "performance_optimization"
    RESEARCH = "research"
    INNOVATION = "innovation"


class CoordinationStrategy(Enum):
    """Strategies for task coordination."""
    SEQUENTIAL = "sequential"  # Tasks executed one after another
    PARALLEL = "parallel"     # Tasks executed simultaneously
    PIPELINE = "pipeline"     # Tasks form a processing pipeline
    ADAPTIVE = "adaptive"     # Strategy chosen based on context
    HIERARCHICAL = "hierarchical"  # Nested task decomposition


@dataclass
class Task:
    """Represents a task in the coordination system."""
    task_id: str
    task_type: TaskType
    description: str
    priority: Priority
    required_capabilities: List[AgentCapability]
    input_data: Dict[str, Any]
    dependencies: List[str] = field(default_factory=list)
    assigned_agents: List[str] = field(default_factory=list)
    status: str = "pending"
    progress: float = 0.0
    created_at: datetime = field(default_factory=datetime.now)
    deadline: Optional[datetime] = None
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None


@dataclass
class CoordinationPlan:
    """Represents a coordination plan for multiple tasks."""
    plan_id: str
    goal: str
    strategy: CoordinationStrategy
    tasks: List[Task]
    resource_allocation: Dict[str, List[str]]
    timeline: Dict[str, datetime]
    success_criteria: Dict[str, float]
    fallback_plans: List[str] = field(default_factory=list)


class CoordinatorAgent(BaseAgent):
    """
    Master Coordinator Agent that orchestrates all other agents in the system.
    Implements intelligent task decomposition, resource allocation, and coordination.
    """
    
    def __init__(
        self,
        agent_id: str = "coordinator",
        message_bus: Optional[MessageBus] = None,
        max_concurrent_tasks: int = 10,
        planning_horizon: timedelta = timedelta(hours=24)
    ):
        super().__init__(
            agent_id=agent_id,
            agent_type="coordinator",
            capabilities=[
                AgentCapability.COORDINATION,
                AgentCapability.PLANNING,
                AgentCapability.MONITORING,
                AgentCapability.OPTIMIZATION
            ],
            message_bus=message_bus
        )
        
        self.logger = get_logger(f"{__name__}.{agent_id}")
        self.max_concurrent_tasks = max_concurrent_tasks
        self.planning_horizon = planning_horizon
        
        # Task and resource management
        self.active_tasks: Dict[str, Task] = {}
        self.completed_tasks: Dict[str, Task] = {}
        self.coordination_plans: Dict[str, CoordinationPlan] = {}
        self.agent_registry: Dict[str, Dict[str, Any]] = {}
        self.resource_usage: Dict[str, float] = {}
        
        # Performance tracking
        self.performance_metrics = {
            "tasks_completed": 0,
            "tasks_failed": 0,
            "average_completion_time": 0.0,
            "resource_efficiency": 0.0,
            "coordination_success_rate": 0.0
        }
        
        # Coordination strategies
        self.strategy_preferences = {
            TaskType.TRAINING: CoordinationStrategy.PIPELINE,
            TaskType.DOCUMENT_PROCESSING: CoordinationStrategy.PARALLEL,
            TaskType.KNOWLEDGE_EXTRACTION: CoordinationStrategy.SEQUENTIAL,
            TaskType.QUALITY_ASSESSMENT: CoordinationStrategy.HIERARCHICAL,
            TaskType.USER_INTERACTION: CoordinationStrategy.ADAPTIVE,
            TaskType.PERFORMANCE_OPTIMIZATION: CoordinationStrategy.ADAPTIVE,
            TaskType.RESEARCH: CoordinationStrategy.PARALLEL,
            TaskType.INNOVATION: CoordinationStrategy.ADAPTIVE
        }
    
    async def initialize(self):
        """Initialize the coordinator agent."""
        await super().initialize()
        self.logger.info("Coordinator Agent initialized")
        
        # Subscribe to relevant message types
        if self.message_bus:
            await self.message_bus.subscribe(MessageType.TASK_REQUEST, self._handle_task_request)
            await self.message_bus.subscribe(MessageType.TASK_COMPLETED, self._handle_task_completion)
            await self.message_bus.subscribe(MessageType.TASK_FAILED, self._handle_task_failure)
            await self.message_bus.subscribe(MessageType.AGENT_STATUS, self._handle_agent_status)
            await self.message_bus.subscribe(MessageType.RESOURCE_REQUEST, self._handle_resource_request)
        
        # Start background coordination tasks
        asyncio.create_task(self._coordination_loop())
        asyncio.create_task(self._performance_monitoring())
        asyncio.create_task(self._resource_optimization())
    
    async def coordinate_goal(
        self,
        goal: str,
        context: Dict[str, Any],
        strategy: Optional[CoordinationStrategy] = None,
        deadline: Optional[datetime] = None
    ) -> str:
        """
        Main entry point for coordinating a high-level goal.
        
        Args:
            goal: High-level goal description
            context: Context and requirements
            strategy: Preferred coordination strategy
            deadline: Deadline for completion
            
        Returns:
            plan_id: Identifier for the coordination plan
        """
        self.logger.info(f"Coordinating goal: {goal}")
        
        try:
            # Analyze goal and decompose into tasks
            tasks = await self._decompose_goal(goal, context)
            
            # Determine coordination strategy
            if not strategy:
                strategy = await self._select_coordination_strategy(tasks, context)
            
            # Create coordination plan
            plan = await self._create_coordination_plan(
                goal, tasks, strategy, deadline
            )
            
            # Allocate resources and assign agents
            await self._allocate_resources(plan)
            
            # Start execution
            await self._execute_coordination_plan(plan)
            
            self.coordination_plans[plan.plan_id] = plan
            self.logger.info(f"Coordination plan {plan.plan_id} started")
            
            return plan.plan_id
            
        except Exception as e:
            self.logger.error(f"Goal coordination failed: {str(e)}")
            raise CoordinationError(f"Failed to coordinate goal: {str(e)}") from e
    
    async def _decompose_goal(
        self,
        goal: str,
        context: Dict[str, Any]
    ) -> List[Task]:
        """Decompose high-level goal into specific tasks."""
        self.logger.debug(f"Decomposing goal: {goal}")
        
        # This is where intelligent goal decomposition would happen
        # For now, implementing basic patterns based on goal type
        
        tasks = []
        
        # Analyze goal to determine task types needed
        goal_lower = goal.lower()
        
        if "train" in goal_lower or "learn" in goal_lower:
            tasks.extend(await self._create_training_tasks(goal, context))
        
        if "process" in goal_lower or "document" in goal_lower:
            tasks.extend(await self._create_document_processing_tasks(goal, context))
        
        if "analyze" in goal_lower or "extract" in goal_lower:
            tasks.extend(await self._create_analysis_tasks(goal, context))
        
        if "improve" in goal_lower or "optimize" in goal_lower:
            tasks.extend(await self._create_optimization_tasks(goal, context))
        
        # If no specific patterns matched, create a general research task
        if not tasks:
            tasks.append(await self._create_general_task(goal, context))
        
        # Add dependencies and priorities
        await self._establish_task_dependencies(tasks)
        
        return tasks
    
    async def _create_training_tasks(
        self,
        goal: str,
        context: Dict[str, Any]
    ) -> List[Task]:
        """Create tasks for training-related goals."""
        tasks = []
        
        # Data preparation task
        tasks.append(Task(
            task_id=str(uuid.uuid4()),
            task_type=TaskType.TRAINING,
            description="Prepare training data and validate quality",
            priority=Priority.HIGH,
            required_capabilities=[
                AgentCapability.DATA_PROCESSING,
                AgentCapability.QUALITY_ASSESSMENT
            ],
            input_data={"goal": goal, "context": context}
        ))
        
        # Model training task
        tasks.append(Task(
            task_id=str(uuid.uuid4()),
            task_type=TaskType.TRAINING,
            description="Execute training process with optimization",
            priority=Priority.HIGH,
            required_capabilities=[
                AgentCapability.TRAINING,
                AgentCapability.OPTIMIZATION
            ],
            input_data={"goal": goal, "context": context}
        ))
        
        # Quality validation task
        tasks.append(Task(
            task_id=str(uuid.uuid4()),
            task_type=TaskType.QUALITY_ASSESSMENT,
            description="Validate training results and performance",
            priority=Priority.MEDIUM,
            required_capabilities=[
                AgentCapability.QUALITY_ASSESSMENT,
                AgentCapability.EVALUATION
            ],
            input_data={"goal": goal, "context": context}
        ))
        
        return tasks
    
    async def _create_document_processing_tasks(
        self,
        goal: str,
        context: Dict[str, Any]
    ) -> List[Task]:
        """Create tasks for document processing goals."""
        tasks = []
        
        # Document analysis task
        tasks.append(Task(
            task_id=str(uuid.uuid4()),
            task_type=TaskType.DOCUMENT_PROCESSING,
            description="Analyze and extract content from documents",
            priority=Priority.HIGH,
            required_capabilities=[
                AgentCapability.DOCUMENT_PROCESSING,
                AgentCapability.EXTRACTION
            ],
            input_data={"goal": goal, "context": context}
        ))
        
        # Knowledge extraction task
        tasks.append(Task(
            task_id=str(uuid.uuid4()),
            task_type=TaskType.KNOWLEDGE_EXTRACTION,
            description="Extract and structure knowledge from content",
            priority=Priority.MEDIUM,
            required_capabilities=[
                AgentCapability.KNOWLEDGE_EXTRACTION,
                AgentCapability.STRUCTURING
            ],
            input_data={"goal": goal, "context": context}
        ))
        
        return tasks
    
    async def _create_analysis_tasks(
        self,
        goal: str,
        context: Dict[str, Any]
    ) -> List[Task]:
        """Create tasks for analysis and extraction goals."""
        tasks = []
        
        tasks.append(Task(
            task_id=str(uuid.uuid4()),
            task_type=TaskType.KNOWLEDGE_EXTRACTION,
            description="Perform comprehensive analysis and extraction",
            priority=Priority.HIGH,
            required_capabilities=[
                AgentCapability.ANALYSIS,
                AgentCapability.EXTRACTION,
                AgentCapability.REASONING
            ],
            input_data={"goal": goal, "context": context}
        ))
        
        return tasks
    
    async def _create_optimization_tasks(
        self,
        goal: str,
        context: Dict[str, Any]
    ) -> List[Task]:
        """Create tasks for optimization goals."""
        tasks = []
        
        tasks.append(Task(
            task_id=str(uuid.uuid4()),
            task_type=TaskType.PERFORMANCE_OPTIMIZATION,
            description="Analyze and optimize system performance",
            priority=Priority.HIGH,
            required_capabilities=[
                AgentCapability.OPTIMIZATION,
                AgentCapability.MONITORING,
                AgentCapability.ANALYSIS
            ],
            input_data={"goal": goal, "context": context}
        ))
        
        return tasks
    
    async def _create_general_task(
        self,
        goal: str,
        context: Dict[str, Any]
    ) -> Task:
        """Create a general research task for unspecified goals."""
        return Task(
            task_id=str(uuid.uuid4()),
            task_type=TaskType.RESEARCH,
            description=f"Research and address: {goal}",
            priority=Priority.MEDIUM,
            required_capabilities=[
                AgentCapability.RESEARCH,
                AgentCapability.ANALYSIS,
                AgentCapability.REASONING
            ],
            input_data={"goal": goal, "context": context}
        )
    
    async def _establish_task_dependencies(self, tasks: List[Task]):
        """Establish dependencies between tasks."""
        # Simple dependency logic - can be enhanced with more sophisticated analysis
        for i, task in enumerate(tasks):
            if i > 0:
                # Each task depends on the previous one by default
                task.dependencies.append(tasks[i-1].task_id)
    
    async def _select_coordination_strategy(
        self,
        tasks: List[Task],
        context: Dict[str, Any]
    ) -> CoordinationStrategy:
        """Select the best coordination strategy for the given tasks."""
        # Analyze task types and determine best strategy
        task_types = [task.task_type for task in tasks]
        
        # If all tasks are the same type, use preference for that type
        if len(set(task_types)) == 1:
            return self.strategy_preferences.get(task_types[0], CoordinationStrategy.ADAPTIVE)
        
        # For mixed task types, consider complexity and dependencies
        if len(tasks) > 5:
            return CoordinationStrategy.HIERARCHICAL
        elif any(task.dependencies for task in tasks):
            return CoordinationStrategy.PIPELINE
        else:
            return CoordinationStrategy.PARALLEL
    
    async def _create_coordination_plan(
        self,
        goal: str,
        tasks: List[Task],
        strategy: CoordinationStrategy,
        deadline: Optional[datetime]
    ) -> CoordinationPlan:
        """Create a detailed coordination plan."""
        plan_id = str(uuid.uuid4())
        
        # Estimate timeline based on task complexity
        timeline = await self._estimate_timeline(tasks, strategy)
        
        # Define success criteria
        success_criteria = {
            "task_completion_rate": 0.95,
            "quality_threshold": 0.8,
            "deadline_adherence": 1.0 if deadline else 0.9
        }
        
        plan = CoordinationPlan(
            plan_id=plan_id,
            goal=goal,
            strategy=strategy,
            tasks=tasks,
            resource_allocation={},
            timeline=timeline,
            success_criteria=success_criteria
        )
        
        return plan
    
    async def _estimate_timeline(
        self,
        tasks: List[Task],
        strategy: CoordinationStrategy
    ) -> Dict[str, datetime]:
        """Estimate timeline for task completion."""
        timeline = {}
        current_time = datetime.now()
        
        # Simple estimation - can be enhanced with historical data
        base_duration = timedelta(minutes=30)  # Base task duration
        
        if strategy == CoordinationStrategy.SEQUENTIAL:
            for i, task in enumerate(tasks):
                start_time = current_time + i * base_duration
                end_time = start_time + base_duration
                timeline[task.task_id] = {
                    "start": start_time,
                    "end": end_time
                }
        elif strategy == CoordinationStrategy.PARALLEL:
            for task in tasks:
                timeline[task.task_id] = {
                    "start": current_time,
                    "end": current_time + base_duration
                }
        else:  # Pipeline or adaptive
            for i, task in enumerate(tasks):
                start_time = current_time + i * base_duration * 0.5
                end_time = start_time + base_duration
                timeline[task.task_id] = {
                    "start": start_time,
                    "end": end_time
                }
        
        return timeline
    
    async def _allocate_resources(self, plan: CoordinationPlan):
        """Allocate agents and resources to tasks."""
        self.logger.debug(f"Allocating resources for plan {plan.plan_id}")
        
        for task in plan.tasks:
            # Find suitable agents for this task
            suitable_agents = await self._find_suitable_agents(task)
            
            # Assign agents based on availability and capability
            assigned = await self._assign_agents(task, suitable_agents)
            
            task.assigned_agents = assigned
            plan.resource_allocation[task.task_id] = assigned
    
    async def _find_suitable_agents(self, task: Task) -> List[str]:
        """Find agents capable of handling the task."""
        suitable_agents = []
        
        for agent_id, agent_info in self.agent_registry.items():
            agent_capabilities = agent_info.get("capabilities", [])
            
            # Check if agent has required capabilities
            if all(cap in agent_capabilities for cap in task.required_capabilities):
                # Check agent availability
                if agent_info.get("status") == AgentStatus.IDLE:
                    suitable_agents.append(agent_id)
        
        return suitable_agents
    
    async def _assign_agents(self, task: Task, suitable_agents: List[str]) -> List[str]:
        """Assign specific agents to a task."""
        # Simple assignment - take first available agent with required capabilities
        # Can be enhanced with load balancing and optimization
        
        if not suitable_agents:
            self.logger.warning(f"No suitable agents found for task {task.task_id}")
            return []
        
        # For now, assign the first suitable agent
        assigned_agent = suitable_agents[0]
        
        # Update agent status
        if assigned_agent in self.agent_registry:
            self.agent_registry[assigned_agent]["status"] = AgentStatus.BUSY
        
        return [assigned_agent]
    
    async def _execute_coordination_plan(self, plan: CoordinationPlan):
        """Execute the coordination plan."""
        self.logger.info(f"Executing coordination plan {plan.plan_id}")
        
        if plan.strategy == CoordinationStrategy.SEQUENTIAL:
            await self._execute_sequential(plan)
        elif plan.strategy == CoordinationStrategy.PARALLEL:
            await self._execute_parallel(plan)
        elif plan.strategy == CoordinationStrategy.PIPELINE:
            await self._execute_pipeline(plan)
        else:
            await self._execute_adaptive(plan)
    
    async def _execute_sequential(self, plan: CoordinationPlan):
        """Execute tasks sequentially."""
        for task in plan.tasks:
            await self._start_task(task)
            # Wait for task completion before starting next
            await self._wait_for_task_completion(task)
    
    async def _execute_parallel(self, plan: CoordinationPlan):
        """Execute tasks in parallel."""
        task_coroutines = [self._start_task(task) for task in plan.tasks]
        await asyncio.gather(*task_coroutines)
    
    async def _execute_pipeline(self, plan: CoordinationPlan):
        """Execute tasks in pipeline fashion."""
        # Start tasks with dependencies consideration
        started_tasks = set()
        
        while len(started_tasks) < len(plan.tasks):
            for task in plan.tasks:
                if task.task_id not in started_tasks:
                    # Check if dependencies are completed
                    dependencies_met = all(
                        dep_id in self.completed_tasks 
                        for dep_id in task.dependencies
                    )
                    
                    if dependencies_met:
                        await self._start_task(task)
                        started_tasks.add(task.task_id)
            
            # Small delay to prevent busy waiting
            await asyncio.sleep(0.1)
    
    async def _execute_adaptive(self, plan: CoordinationPlan):
        """Execute tasks using adaptive strategy."""
        # Implement adaptive execution based on current conditions
        # For now, default to pipeline execution
        await self._execute_pipeline(plan)
    
    async def _start_task(self, task: Task):
        """Start execution of a specific task."""
        self.logger.debug(f"Starting task {task.task_id}")
        
        task.status = "running"
        self.active_tasks[task.task_id] = task
        
        # Send task to assigned agents
        for agent_id in task.assigned_agents:
            message = AgentMessage(
                message_id=str(uuid.uuid4()),
                sender_id=self.agent_id,
                recipient_id=agent_id,
                message_type=MessageType.TASK_ASSIGNMENT,
                content={
                    "task": task.__dict__,
                    "deadline": task.deadline.isoformat() if task.deadline else None
                },
                priority=task.priority
            )
            
            if self.message_bus:
                await self.message_bus.send_message(message)
    
    async def _wait_for_task_completion(self, task: Task):
        """Wait for a task to complete."""
        while task.status == "running":
            await asyncio.sleep(1.0)
    
    # Message handlers
    async def _handle_task_request(self, message: AgentMessage):
        """Handle incoming task requests."""
        self.logger.debug(f"Received task request: {message.message_id}")
        
        # Extract task information from message
        task_info = message.content
        
        # Create and coordinate the task
        await self.coordinate_goal(
            goal=task_info.get("goal", ""),
            context=task_info.get("context", {}),
            deadline=task_info.get("deadline")
        )
    
    async def _handle_task_completion(self, message: AgentMessage):
        """Handle task completion notifications."""
        task_id = message.content.get("task_id")
        result = message.content.get("result")
        
        if task_id in self.active_tasks:
            task = self.active_tasks[task_id]
            task.status = "completed"
            task.result = result
            task.progress = 1.0
            
            # Move to completed tasks
            self.completed_tasks[task_id] = task
            del self.active_tasks[task_id]
            
            # Update performance metrics
            self.performance_metrics["tasks_completed"] += 1
            
            # Release assigned agents
            for agent_id in task.assigned_agents:
                if agent_id in self.agent_registry:
                    self.agent_registry[agent_id]["status"] = AgentStatus.IDLE
            
            self.logger.info(f"Task {task_id} completed successfully")
    
    async def _handle_task_failure(self, message: AgentMessage):
        """Handle task failure notifications."""
        task_id = message.content.get("task_id")
        error = message.content.get("error")
        
        if task_id in self.active_tasks:
            task = self.active_tasks[task_id]
            task.status = "failed"
            task.error = error
            
            # Move to completed tasks (as failed)
            self.completed_tasks[task_id] = task
            del self.active_tasks[task_id]
            
            # Update performance metrics
            self.performance_metrics["tasks_failed"] += 1
            
            # Release assigned agents
            for agent_id in task.assigned_agents:
                if agent_id in self.agent_registry:
                    self.agent_registry[agent_id]["status"] = AgentStatus.IDLE
            
            self.logger.error(f"Task {task_id} failed: {error}")
            
            # Implement failure recovery if needed
            await self._handle_task_failure_recovery(task)
    
    async def _handle_task_failure_recovery(self, failed_task: Task):
        """Handle recovery from task failure."""
        # Implement retry logic or alternative approaches
        self.logger.info(f"Attempting recovery for failed task {failed_task.task_id}")
        
        # Simple retry logic - can be enhanced
        if failed_task.input_data.get("retry_count", 0) < 3:
            retry_task = Task(
                task_id=str(uuid.uuid4()),
                task_type=failed_task.task_type,
                description=f"Retry: {failed_task.description}",
                priority=failed_task.priority,
                required_capabilities=failed_task.required_capabilities,
                input_data={
                    **failed_task.input_data,
                    "retry_count": failed_task.input_data.get("retry_count", 0) + 1
                },
                dependencies=failed_task.dependencies
            )
            
            await self._start_task(retry_task)
    
    async def _handle_agent_status(self, message: AgentMessage):
        """Handle agent status updates."""
        agent_id = message.sender_id
        status_info = message.content
        
        # Update agent registry
        if agent_id not in self.agent_registry:
            self.agent_registry[agent_id] = {}
        
        self.agent_registry[agent_id].update(status_info)
        
        self.logger.debug(f"Updated status for agent {agent_id}")
    
    async def _handle_resource_request(self, message: AgentMessage):
        """Handle resource allocation requests."""
        self.logger.debug(f"Received resource request from {message.sender_id}")
        
        # Implement resource allocation logic
        requested_resources = message.content.get("resources", [])
        available_resources = await self._check_resource_availability(requested_resources)
        
        # Respond with resource allocation
        response = AgentMessage(
            message_id=str(uuid.uuid4()),
            sender_id=self.agent_id,
            recipient_id=message.sender_id,
            message_type=MessageType.RESOURCE_ALLOCATION,
            content={"allocated_resources": available_resources},
            priority=Priority.MEDIUM
        )
        
        if self.message_bus:
            await self.message_bus.send_message(response)
    
    async def _check_resource_availability(self, requested_resources: List[str]) -> List[str]:
        """Check availability of requested resources."""
        # Simple implementation - can be enhanced with actual resource tracking
        return requested_resources  # For now, assume all resources are available
    
    # Background processes
    async def _coordination_loop(self):
        """Main coordination loop."""
        self.logger.debug("Starting coordination loop")
        
        while self.status == AgentStatus.IDLE:  # Use IDLE instead of ACTIVE
            try:
                # Monitor active tasks
                await self._monitor_active_tasks()
                
                # Check for plan completion
                await self._check_plan_completion()
                
                # Optimize current allocations
                await self._optimize_current_allocations()
                
                # Brief sleep to prevent excessive CPU usage
                await asyncio.sleep(1.0)
                
            except Exception as e:
                self.logger.error(f"Coordination loop error: {str(e)}")
                await asyncio.sleep(5.0)  # Wait longer on error
    
    async def _monitor_active_tasks(self):
        """Monitor the progress of active tasks."""
        for plan_id, plan in list(self.coordination_plans.items()):
            for task in plan.tasks:
                if task.status == "running":
                    # Check for task timeout or issues
                    pass
    
    async def _check_plan_completion(self):
        """Check if any plans have completed."""
        completed_plans = []
        
        for plan_id, plan in self.coordination_plans.items():
            if all(task.status in ["completed", "failed"] for task in plan.tasks):
                completed_plans.append(plan_id)
        
        for plan_id in completed_plans:
            plan = self.coordination_plans.pop(plan_id)
            await self._finalize_plan(plan)
    
    async def _finalize_plan(self, plan: CoordinationPlan):
        """Finalize a completed coordination plan."""
        # Calculate plan success
        completed_tasks = [t for t in plan.tasks if t.status == "completed"]
        success_rate = len(completed_tasks) / len(plan.tasks) if plan.tasks else 0.0
        
        self.completed_tasks[plan.plan_id] = {
            "plan": plan,
            "completed_at": datetime.now(),
            "success_rate": success_rate,
            "total_tasks": len(plan.tasks),
            "completed_tasks": len(completed_tasks)
        }
        
        self.logger.info(
            f"Plan {plan.plan_id} completed with {success_rate:.1%} success rate "
            f"({len(completed_tasks)}/{len(plan.tasks)} tasks)"
        )
    
    async def _optimize_current_allocations(self):
        """Optimize current resource allocations."""
        # This is a placeholder for optimization logic
        pass
    
    async def _performance_monitoring(self):
        """Monitor performance metrics."""
        while self.status == AgentStatus.IDLE:  # Use IDLE instead of ACTIVE
            try:
                await self._update_performance_metrics()
                await asyncio.sleep(30.0)  # Update every 30 seconds
            except Exception as e:
                self.logger.error(f"Performance monitoring error: {str(e)}")
                await asyncio.sleep(60.0)
    
    async def _update_performance_metrics(self):
        """Update performance metrics."""
        current_time = datetime.now()
        
        # Update coordination metrics
        total_tasks = sum(len(plan.tasks) for plan in self.coordination_plans.values())
        completed_tasks = sum(
            len([t for t in plan.tasks if t.status == "completed"])
            for plan in self.coordination_plans.values()
        )
        
        self.performance_metrics.update({
            "active_plans": len(self.coordination_plans),
            "completed_plans": len(self.completed_tasks),
            "total_active_tasks": total_tasks,
            "completed_tasks": completed_tasks,
            "available_agents": len(self.agent_registry),
            "last_updated": current_time.isoformat()
        })
        
        # Log metrics periodically
        if hasattr(self, '_last_metrics_log'):
            if (current_time - self._last_metrics_log).seconds > 300:  # Every 5 minutes
                self._log_performance_metrics()
                self._last_metrics_log = current_time
        else:
            self._last_metrics_log = current_time
    
    def _log_performance_metrics(self):
        """Log current performance metrics."""
        metrics = self.performance_metrics
        self.logger.info(
            f"Coordination metrics: Plans={metrics.get('active_plans', 0)}, "
            f"Tasks={metrics.get('total_active_tasks', 0)}, "
            f"Agents={metrics.get('available_agents', 0)}"
        )
    
    async def _resource_optimization(self):
        """Optimize resource allocation."""
        while self.status == AgentStatus.IDLE:  # Use IDLE instead of ACTIVE
            try:
                await self._rebalance_workloads()
                await self._scale_resources()
                await asyncio.sleep(60.0)  # Optimize every minute
            except Exception as e:
                self.logger.error(f"Resource optimization error: {str(e)}")
                await asyncio.sleep(120.0)
    
    async def _rebalance_workloads(self):
        """Rebalance workloads across agents."""
        # Placeholder for workload rebalancing logic
        pass
    
    async def _scale_resources(self):
        """Scale resources based on demand."""
        # Placeholder for resource scaling logic
        pass
    
    # Public API methods
    def get_status(self) -> Dict[str, Any]:
        """Get current coordinator status."""
        return {
            "agent_id": self.agent_id,
            "status": self.status.value,
            "active_tasks": len(self.active_tasks),
            "completed_tasks": len(self.completed_tasks),
            "active_plans": len(self.coordination_plans),
            "registered_agents": len(self.agent_registry),
            "performance_metrics": self.performance_metrics
        }
    
    def get_task_status(self, task_id: str) -> Optional[Dict[str, Any]]:
        """Get status of a specific task."""
        if task_id in self.active_tasks:
            return self.active_tasks[task_id].__dict__
        elif task_id in self.completed_tasks:
            return self.completed_tasks[task_id].__dict__
        else:
            return None
    
    def get_plan_status(self, plan_id: str) -> Optional[Dict[str, Any]]:
        """Get status of a coordination plan."""
        if plan_id in self.coordination_plans:
            plan = self.coordination_plans[plan_id]
            return {
                "plan_id": plan.plan_id,
                "goal": plan.goal,
                "strategy": plan.strategy.value,
                "total_tasks": len(plan.tasks),
                "completed_tasks": sum(
                    1 for task in plan.tasks
                    if task.task_id in self.completed_tasks
                ),
                "progress": self._calculate_plan_progress(plan)
            }
        else:
            return None
    
    def _calculate_plan_progress(self, plan: CoordinationPlan) -> float:
        """Calculate overall progress of a coordination plan."""
        if not plan.tasks:
            return 1.0
        
        total_progress = sum(
            self.completed_tasks.get(task.task_id, task).progress
            for task in plan.tasks
        )
        
        return total_progress / len(plan.tasks)
    
    async def process_task(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Process a task assigned to this coordinator agent.
        
        Args:
            task_data: Task information and parameters
            
        Returns:
            Task result or error information
        """
        try:
            task_type = task_data.get("type", "coordination")
            
            if task_type == "coordination":
                # Handle coordination task
                goal = task_data.get("goal", "")
                context = task_data.get("context", {})
                strategy = task_data.get("strategy")
                deadline = task_data.get("deadline")
                
                # Convert deadline string to datetime if provided
                if deadline and isinstance(deadline, str):
                    deadline = datetime.fromisoformat(deadline)
                
                # Coordinate the goal
                plan_id = await self.coordinate_goal(goal, context, strategy, deadline)
                
                return {
                    "status": "success",
                    "plan_id": plan_id,
                    "message": f"Coordination plan {plan_id} created successfully"
                }
            
            elif task_type == "status_check":
                # Handle status check task
                return {
                    "status": "success",
                    "coordinator_status": self.get_status()
                }
            
            elif task_type == "plan_status":
                # Handle plan status check
                plan_id = task_data.get("plan_id")
                plan_status = self.get_plan_status(plan_id)
                
                if plan_status:
                    return {
                        "status": "success",
                        "plan_status": plan_status
                    }
                else:
                    return {
                        "status": "error",
                        "message": f"Plan {plan_id} not found"
                    }
            
            else:
                return {
                    "status": "error",
                    "message": f"Unknown task type: {task_type}"
                }
                
        except Exception as e:
            self.logger.error(f"Error processing task: {str(e)}")
            return {
                "status": "error",
                "message": str(e)
            } 