"""
Base agent classes and interfaces for the Toxo agent system.

This module provides the foundational classes and enums that all agents
in the Toxo platform inherit from or implement.
"""

import asyncio
import time
import uuid
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional, Set

from ..utils.logger import get_logger
from .communication import MessageBus, AgentMessage, MessageType, Priority


class AgentStatus(Enum):
    """Status of an agent."""
    INITIALIZING = "initializing"
    IDLE = "idle"
    BUSY = "busy"
    ERROR = "error"
    OFFLINE = "offline"
    MAINTENANCE = "maintenance"


class AgentCapability(Enum):
    """Capabilities that agents can have."""
    # Core capabilities
    COORDINATION = "coordination"
    PLANNING = "planning"
    MONITORING = "monitoring"
    OPTIMIZATION = "optimization"
    
    # Training capabilities
    TRAINING = "training"
    VALIDATION = "validation"
    EVALUATION = "evaluation"
    
    # Data processing capabilities
    DOCUMENT_PROCESSING = "document_processing"
    TEXT_ANALYSIS = "text_analysis"
    DATA_EXTRACTION = "data_extraction"
    
    # Knowledge capabilities
    KNOWLEDGE_EXTRACTION = "knowledge_extraction"
    KNOWLEDGE_SYNTHESIS = "knowledge_synthesis"
    QUESTION_ANSWERING = "question_answering"
    
    # Quality capabilities
    QUALITY_ASSESSMENT = "quality_assessment"
    ERROR_DETECTION = "error_detection"
    PERFORMANCE_ANALYSIS = "performance_analysis"
    
    # User interaction capabilities
    USER_INTERACTION = "user_interaction"
    FEEDBACK_PROCESSING = "feedback_processing"
    EXPLANATION_GENERATION = "explanation_generation"
    
    # Research capabilities
    RESEARCH = "research"
    INNOVATION = "innovation"
    EXPERIMENTATION = "experimentation"


@dataclass
class AgentMetrics:
    """Metrics for tracking agent performance."""
    tasks_completed: int = 0
    tasks_failed: int = 0
    average_response_time: float = 0.0
    cpu_usage: float = 0.0
    memory_usage: float = 0.0
    uptime: float = 0.0
    last_activity: Optional[datetime] = None
    error_rate: float = 0.0
    throughput: float = 0.0  # Tasks per minute
    
    def update_completion(self, response_time: float):
        """Update metrics for a completed task."""
        self.tasks_completed += 1
        total_tasks = self.tasks_completed + self.tasks_failed
        if total_tasks > 0:
            self.average_response_time = (
                (self.average_response_time * (total_tasks - 1) + response_time) / total_tasks
            )
        self.last_activity = datetime.now()
        self._update_error_rate()
    
    def update_failure(self):
        """Update metrics for a failed task."""
        self.tasks_failed += 1
        self.last_activity = datetime.now()
        self._update_error_rate()
    
    def _update_error_rate(self):
        """Calculate error rate."""
        total_tasks = self.tasks_completed + self.tasks_failed
        if total_tasks > 0:
            self.error_rate = self.tasks_failed / total_tasks
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert metrics to dictionary."""
        return {
            "tasks_completed": self.tasks_completed,
            "tasks_failed": self.tasks_failed,
            "average_response_time": self.average_response_time,
            "cpu_usage": self.cpu_usage,
            "memory_usage": self.memory_usage,
            "uptime": self.uptime,
            "last_activity": self.last_activity.isoformat() if self.last_activity else None,
            "error_rate": self.error_rate,
            "throughput": self.throughput
        }


class BaseAgent(ABC):
    """
    Base class for all agents in the Toxo system.
    
    Provides common functionality for agent lifecycle, communication,
    and monitoring.
    """
    
    def __init__(
        self,
        agent_id: str,
        agent_type: str,
        capabilities: List[AgentCapability],
        message_bus: Optional[MessageBus] = None
    ):
        self.agent_id = agent_id
        self.agent_type = agent_type
        self.capabilities = set(capabilities)
        self.message_bus = message_bus
        
        # Status and metrics
        self.status = AgentStatus.INITIALIZING
        self.metrics = AgentMetrics()
        self.start_time = datetime.now()
        
        # Configuration
        self.config: Dict[str, Any] = {}
        self.logger = get_logger(f"{__name__}.{agent_id}")
        
        # Task management
        self.current_tasks: Set[str] = set()
        self.task_history: List[Dict[str, Any]] = []
        
        # Health monitoring
        self.heartbeat_interval = 30.0  # seconds
        self.last_heartbeat = datetime.now()
        
        self.logger.info(f"Agent {agent_id} ({agent_type}) initialized")
    
    async def initialize(self):
        """Initialize the agent."""
        self.status = AgentStatus.IDLE
        self.logger.info(f"Agent {self.agent_id} initialized")
        
        # Start heartbeat if message bus is available
        if self.message_bus:
            asyncio.create_task(self._heartbeat_loop())
    
    async def shutdown(self):
        """Shutdown the agent gracefully."""
        self.status = AgentStatus.OFFLINE
        self.logger.info(f"Agent {self.agent_id} shutting down")
        
        # Send shutdown notification
        if self.message_bus:
            message = AgentMessage(
                message_type=MessageType.SHUTDOWN,
                sender_id=self.agent_id,
                payload={"reason": "graceful_shutdown"}
            )
            await self.message_bus.publish(message)
    
    @abstractmethod
    async def process_task(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Process a task assigned to this agent.
        
        Args:
            task_data: Task information and parameters
            
        Returns:
            Task result or error information
        """
        pass
    
    async def handle_message(self, message: AgentMessage):
        """
        Handle incoming messages.
        
        Args:
            message: Message to handle
        """
        try:
            if message.message_type == MessageType.TASK_REQUEST:
                await self._handle_task_request(message)
            elif message.message_type == MessageType.HEARTBEAT:
                await self._handle_heartbeat(message)
            else:
                await self._handle_custom_message(message)
        except Exception as e:
            self.logger.error(f"Error handling message {message.message_id}: {str(e)}")
            self.metrics.update_failure()
    
    async def _handle_task_request(self, message: AgentMessage):
        """Handle a task request message."""
        task_id = message.payload.get("task_id")
        if not task_id:
            self.logger.warning(f"Task request {message.message_id} missing task_id")
            return
        
        if task_id in self.current_tasks:
            self.logger.warning(f"Task {task_id} already in progress")
            return
        
        self.current_tasks.add(task_id)
        self.status = AgentStatus.BUSY
        start_time = time.time()
        
        try:
            # Process the task
            result = await self.process_task(message.payload)
            
            # Calculate response time
            response_time = time.time() - start_time
            self.metrics.update_completion(response_time)
            
            # Send completion message
            if self.message_bus:
                response = AgentMessage(
                    message_type=MessageType.TASK_COMPLETED,
                    sender_id=self.agent_id,
                    recipient_id=message.sender_id,
                    correlation_id=message.correlation_id,
                    payload={
                        "task_id": task_id,
                        "result": result,
                        "response_time": response_time
                    }
                )
                await self.message_bus.send_response(response)
            
            # Update task history
            self.task_history.append({
                "task_id": task_id,
                "status": "completed",
                "start_time": start_time,
                "completion_time": time.time(),
                "response_time": response_time
            })
            
        except Exception as e:
            self.logger.error(f"Task {task_id} failed: {str(e)}")
            self.metrics.update_failure()
            
            # Send failure message
            if self.message_bus:
                response = AgentMessage(
                    message_type=MessageType.TASK_FAILED,
                    sender_id=self.agent_id,
                    recipient_id=message.sender_id,
                    correlation_id=message.correlation_id,
                    payload={
                        "task_id": task_id,
                        "error": str(e),
                        "response_time": time.time() - start_time
                    }
                )
                await self.message_bus.send_response(response)
            
            # Update task history
            self.task_history.append({
                "task_id": task_id,
                "status": "failed",
                "start_time": start_time,
                "completion_time": time.time(),
                "error": str(e)
            })
        
        finally:
            self.current_tasks.remove(task_id)
            if not self.current_tasks:
                self.status = AgentStatus.IDLE
    
    async def _handle_heartbeat(self, message: AgentMessage):
        """Handle a heartbeat message."""
        self.last_heartbeat = datetime.now()
        
        # Respond with status if requested
        if message.payload.get("request_status"):
            status_message = AgentMessage(
                message_type=MessageType.AGENT_STATUS,
                sender_id=self.agent_id,
                recipient_id=message.sender_id,
                payload=self.get_status()
            )
            if self.message_bus:
                await self.message_bus.publish(status_message)
    
    async def _handle_custom_message(self, message: AgentMessage):
        """Handle custom message types. Override in subclasses."""
        self.logger.debug(f"Received custom message: {message.message_type.value}")
    
    async def _heartbeat_loop(self):
        """Send periodic heartbeat messages."""
        while self.status != AgentStatus.OFFLINE:
            try:
                if self.message_bus:
                    heartbeat = AgentMessage(
                        message_type=MessageType.HEARTBEAT,
                        sender_id=self.agent_id,
                        payload={
                            "status": self.status.value,
                            "uptime": (datetime.now() - self.start_time).total_seconds(),
                            "active_tasks": len(self.current_tasks)
                        }
                    )
                    await self.message_bus.publish(heartbeat)
                
                await asyncio.sleep(self.heartbeat_interval)
                
            except Exception as e:
                self.logger.error(f"Error in heartbeat loop: {str(e)}")
                await asyncio.sleep(self.heartbeat_interval)
    
    def has_capability(self, capability: AgentCapability) -> bool:
        """Check if agent has a specific capability."""
        return capability in self.capabilities
    
    def add_capability(self, capability: AgentCapability):
        """Add a capability to the agent."""
        self.capabilities.add(capability)
        self.logger.debug(f"Added capability: {capability.value}")
    
    def remove_capability(self, capability: AgentCapability):
        """Remove a capability from the agent."""
        self.capabilities.discard(capability)
        self.logger.debug(f"Removed capability: {capability.value}")
    
    def get_status(self) -> Dict[str, Any]:
        """Get current agent status and metrics."""
        uptime = (datetime.now() - self.start_time).total_seconds()
        self.metrics.uptime = uptime
        
        return {
            "agent_id": self.agent_id,
            "agent_type": self.agent_type,
            "status": self.status.value,
            "capabilities": [cap.value for cap in self.capabilities],
            "metrics": self.metrics.to_dict(),
            "current_tasks": list(self.current_tasks),
            "uptime": uptime,
            "last_heartbeat": self.last_heartbeat.isoformat()
        }
    
    def configure(self, config: Dict[str, Any]):
        """Update agent configuration."""
        self.config.update(config)
        self.logger.info(f"Agent {self.agent_id} configured")
    
    def get_task_history(self, limit: int = 100) -> List[Dict[str, Any]]:
        """Get recent task history."""
        return self.task_history[-limit:]
    
    def clear_task_history(self):
        """Clear task history."""
        self.task_history.clear()
        self.logger.debug("Task history cleared") 