"""
Agent communication system for Toxo platform.

This module provides message passing, event handling, and coordination
communication between agents in the system.
"""

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, Callable, Dict, List, Optional, Set
from collections import defaultdict

from ..utils.logger import get_logger


class MessageType(Enum):
    """Types of messages that can be sent between agents."""
    TASK_REQUEST = "task_request"
    TASK_COMPLETED = "task_completed"
    TASK_FAILED = "task_failed"
    TASK_UPDATE = "task_update"
    AGENT_STATUS = "agent_status"
    RESOURCE_REQUEST = "resource_request"
    RESOURCE_RESPONSE = "resource_response"
    COORDINATION_REQUEST = "coordination_request"
    COORDINATION_RESPONSE = "coordination_response"
    ERROR_NOTIFICATION = "error_notification"
    HEARTBEAT = "heartbeat"
    SHUTDOWN = "shutdown"


class Priority(Enum):
    """Message priority levels."""
    LOW = 1
    MEDIUM = 2
    HIGH = 3
    CRITICAL = 4
    URGENT = 5


@dataclass
class AgentMessage:
    """Represents a message between agents."""
    message_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    message_type: MessageType = MessageType.TASK_REQUEST
    sender_id: str = ""
    recipient_id: Optional[str] = None  # None for broadcast
    priority: Priority = Priority.MEDIUM
    timestamp: datetime = field(default_factory=datetime.now)
    ttl: Optional[int] = None  # Time to live in seconds
    payload: Dict[str, Any] = field(default_factory=dict)
    correlation_id: Optional[str] = None  # For request-response patterns
    reply_to: Optional[str] = None
    
    def is_expired(self) -> bool:
        """Check if message has expired based on TTL."""
        if self.ttl is None:
            return False
        elapsed = (datetime.now() - self.timestamp).total_seconds()
        return elapsed > self.ttl
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert message to dictionary representation."""
        return {
            "message_id": self.message_id,
            "message_type": self.message_type.value,
            "sender_id": self.sender_id,
            "recipient_id": self.recipient_id,
            "priority": self.priority.value,
            "timestamp": self.timestamp.isoformat(),
            "ttl": self.ttl,
            "payload": self.payload,
            "correlation_id": self.correlation_id,
            "reply_to": self.reply_to
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'AgentMessage':
        """Create message from dictionary representation."""
        return cls(
            message_id=data.get("message_id", str(uuid.uuid4())),
            message_type=MessageType(data.get("message_type", "task_request")),
            sender_id=data.get("sender_id", ""),
            recipient_id=data.get("recipient_id"),
            priority=Priority(data.get("priority", 2)),
            timestamp=datetime.fromisoformat(data.get("timestamp", datetime.now().isoformat())),
            ttl=data.get("ttl"),
            payload=data.get("payload", {}),
            correlation_id=data.get("correlation_id"),
            reply_to=data.get("reply_to")
        )


class MessageHandler:
    """Handler for processing messages of specific types."""
    
    def __init__(self, handler_func: Callable[[AgentMessage], None], agent_id: str):
        self.handler_func = handler_func
        self.agent_id = agent_id
        self.message_count = 0
        self.last_processed = None
    
    async def handle(self, message: AgentMessage):
        """Handle a message."""
        try:
            await self.handler_func(message)
            self.message_count += 1
            self.last_processed = datetime.now()
        except Exception as e:
            logger = get_logger(f"MessageHandler.{self.agent_id}")
            logger.error(f"Error handling message {message.message_id}: {str(e)}")


class MessageBus:
    """
    Central message bus for agent communication.
    
    Provides publish-subscribe messaging, routing, and delivery guarantees.
    """
    
    def __init__(self, max_queue_size: int = 1000):
        self.logger = get_logger(f"{__name__}.MessageBus")
        self.max_queue_size = max_queue_size
        
        # Message routing and subscriptions
        self.subscribers: Dict[MessageType, List[MessageHandler]] = defaultdict(list)
        self.agent_queues: Dict[str, asyncio.Queue] = {}
        self.broadcast_queue: asyncio.Queue = asyncio.Queue(maxsize=max_queue_size)
        
        # Message tracking
        self.message_history: List[AgentMessage] = []
        self.pending_responses: Dict[str, AgentMessage] = {}
        self.delivery_receipts: Dict[str, bool] = {}
        
        # Bus status
        self.is_running = False
        self.stats = {
            "messages_sent": 0,
            "messages_delivered": 0,
            "messages_failed": 0,
            "active_subscriptions": 0
        }
    
    async def start(self):
        """Start the message bus."""
        self.is_running = True
        self.logger.info("Message bus started")
        
        # Start background tasks
        asyncio.create_task(self._message_router())
        asyncio.create_task(self._cleanup_expired_messages())
    
    async def stop(self):
        """Stop the message bus."""
        self.is_running = False
        self.logger.info("Message bus stopped")
    
    async def publish(self, message: AgentMessage) -> bool:
        """
        Publish a message to the bus.
        
        Args:
            message: Message to publish
            
        Returns:
            bool: True if message was queued successfully
        """
        if message.is_expired():
            self.logger.warning(f"Dropping expired message {message.message_id}")
            return False
        
        try:
            if message.recipient_id:
                # Direct message to specific agent
                if message.recipient_id not in self.agent_queues:
                    self.agent_queues[message.recipient_id] = asyncio.Queue(maxsize=self.max_queue_size)
                
                await self.agent_queues[message.recipient_id].put(message)
            else:
                # Broadcast message
                await self.broadcast_queue.put(message)
            
            self.message_history.append(message)
            self.stats["messages_sent"] += 1
            
            self.logger.debug(f"Published message {message.message_id} from {message.sender_id}")
            return True
            
        except asyncio.QueueFull:
            self.logger.error(f"Queue full, dropping message {message.message_id}")
            self.stats["messages_failed"] += 1
            return False
        except Exception as e:
            self.logger.error(f"Failed to publish message {message.message_id}: {str(e)}")
            self.stats["messages_failed"] += 1
            return False
    
    async def subscribe(
        self,
        message_type: MessageType,
        handler: Callable[[AgentMessage], None],
        agent_id: str = "unknown"
    ):
        """
        Subscribe to messages of a specific type.
        
        Args:
            message_type: Type of messages to subscribe to
            handler: Function to handle messages
            agent_id: ID of the subscribing agent
        """
        message_handler = MessageHandler(handler, agent_id)
        self.subscribers[message_type].append(message_handler)
        self.stats["active_subscriptions"] += 1
        
        self.logger.debug(f"Agent {agent_id} subscribed to {message_type.value}")
    
    async def unsubscribe(self, message_type: MessageType, agent_id: str):
        """
        Unsubscribe from messages of a specific type.
        
        Args:
            message_type: Type of messages to unsubscribe from
            agent_id: ID of the unsubscribing agent
        """
        self.subscribers[message_type] = [
            handler for handler in self.subscribers[message_type]
            if handler.agent_id != agent_id
        ]
        self.stats["active_subscriptions"] -= 1
        
        self.logger.debug(f"Agent {agent_id} unsubscribed from {message_type.value}")
    
    async def request_response(
        self,
        message: AgentMessage,
        timeout: float = 30.0
    ) -> Optional[AgentMessage]:
        """
        Send a request and wait for response.
        
        Args:
            message: Request message
            timeout: Timeout in seconds
            
        Returns:
            Response message or None if timeout
        """
        correlation_id = str(uuid.uuid4())
        message.correlation_id = correlation_id
        
        # Create response future
        response_future = asyncio.Future()
        self.pending_responses[correlation_id] = response_future
        
        # Send request
        success = await self.publish(message)
        if not success:
            del self.pending_responses[correlation_id]
            return None
        
        try:
            # Wait for response
            response = await asyncio.wait_for(response_future, timeout=timeout)
            return response
        except asyncio.TimeoutError:
            self.logger.warning(f"Request {correlation_id} timed out")
            return None
        finally:
            if correlation_id in self.pending_responses:
                del self.pending_responses[correlation_id]
    
    async def send_response(self, response: AgentMessage):
        """
        Send a response to a previous request.
        
        Args:
            response: Response message
        """
        if response.correlation_id and response.correlation_id in self.pending_responses:
            future = self.pending_responses[response.correlation_id]
            if not future.done():
                future.set_result(response)
        else:
            # Regular publish if not a response
            await self.publish(response)
    
    async def _message_router(self):
        """Background task to route messages to subscribers."""
        while self.is_running:
            try:
                # Process broadcast messages
                try:
                    message = await asyncio.wait_for(self.broadcast_queue.get(), timeout=0.1)
                    await self._deliver_to_subscribers(message)
                except asyncio.TimeoutError:
                    pass
                
                # Process direct messages
                for agent_id, queue in self.agent_queues.items():
                    try:
                        message = await asyncio.wait_for(queue.get(), timeout=0.1)
                        await self._deliver_to_subscribers(message)
                    except asyncio.TimeoutError:
                        continue
                
                await asyncio.sleep(0.01)  # Small delay to prevent busy waiting
                
            except Exception as e:
                self.logger.error(f"Error in message router: {str(e)}")
                await asyncio.sleep(1)
    
    async def _deliver_to_subscribers(self, message: AgentMessage):
        """Deliver message to all subscribers of its type."""
        handlers = self.subscribers.get(message.message_type, [])
        
        delivery_tasks = []
        for handler in handlers:
            # Skip if message is for specific recipient and handler is not the recipient
            if message.recipient_id and handler.agent_id != message.recipient_id:
                continue
            
            task = asyncio.create_task(handler.handle(message))
            delivery_tasks.append(task)
        
        if delivery_tasks:
            await asyncio.gather(*delivery_tasks, return_exceptions=True)
            self.stats["messages_delivered"] += 1
    
    async def _cleanup_expired_messages(self):
        """Background task to clean up expired messages."""
        while self.is_running:
            try:
                current_time = datetime.now()
                
                # Clean up message history (keep last 1000 messages)
                if len(self.message_history) > 1000:
                    self.message_history = self.message_history[-1000:]
                
                # Clean up expired pending responses
                expired_correlations = []
                for correlation_id, future in self.pending_responses.items():
                    if future.done():
                        expired_correlations.append(correlation_id)
                
                for correlation_id in expired_correlations:
                    del self.pending_responses[correlation_id]
                
                await asyncio.sleep(60)  # Cleanup every minute
                
            except Exception as e:
                self.logger.error(f"Error in cleanup task: {str(e)}")
                await asyncio.sleep(60)
    
    def get_stats(self) -> Dict[str, Any]:
        """Get message bus statistics."""
        return {
            **self.stats,
            "queue_sizes": {
                agent_id: queue.qsize()
                for agent_id, queue in self.agent_queues.items()
            },
            "broadcast_queue_size": self.broadcast_queue.qsize(),
            "pending_responses": len(self.pending_responses),
            "message_history_size": len(self.message_history),
            "is_running": self.is_running
        }
    
    def get_state(self) -> Dict[str, Any]:
        """Get the current state for serialization."""
        return {
            "registered_agents_count": len(getattr(self, 'agent_queues', {})),
            "pending_responses_count": len(getattr(self, 'pending_responses', {})),
            "message_history_size": len(getattr(self, 'message_history', [])),
            "is_running": getattr(self, 'is_running', False),
            "system_type": "MessageBus"
        } 