"""Protocolo A2A (Agent-to-Agent) para comunicação entre agentes Mangaba"""

import json
import uuid
import threading
from datetime import datetime
from typing import Dict, List, Optional, Any, Callable
from dataclasses import dataclass, asdict
from enum import Enum

class MessageType(Enum):
    """Tipos de mensagens A2A"""
    REQUEST = "request"
    RESPONSE = "response"
    BROADCAST = "broadcast"
    NOTIFICATION = "notification"
    ERROR = "error"

@dataclass
class A2AMessage:
    """Mensagem padrão do protocolo A2A"""
    id: str
    sender_id: str
    receiver_id: Optional[str]
    message_type: MessageType
    content: Dict[str, Any]
    timestamp: str
    correlation_id: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    
    @classmethod
    def create(cls, sender_id: str, message_type: MessageType, content: Dict[str, Any], 
               receiver_id: Optional[str] = None, correlation_id: Optional[str] = None) -> 'A2AMessage':
        """Cria uma nova mensagem A2A"""
        return cls(
            id=str(uuid.uuid4()),
            sender_id=sender_id,
            receiver_id=receiver_id,
            message_type=message_type,
            content=content,
            timestamp=datetime.now().isoformat(),
            correlation_id=correlation_id,
            metadata={}
        )
    
    def to_dict(self) -> Dict[str, Any]:
        """Converte mensagem para dicionário"""
        data = asdict(self)
        data['message_type'] = self.message_type.value
        return data
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'A2AMessage':
        """Cria mensagem a partir de dicionário"""
        data['message_type'] = MessageType(data['message_type'])
        return cls(**data)

class A2AProtocol:
    """Protocolo de comunicação Agent-to-Agent"""
    
    def __init__(self, agent_id: str):
        self.agent_id = agent_id
        self.message_handlers: Dict[MessageType, List[Callable]] = {
            msg_type: [] for msg_type in MessageType
        }
        self.connected_agents: Dict[str, 'A2AAgent'] = {}
        self.message_history: List[A2AMessage] = []
        self._lock = threading.RLock()  # Lock para operações thread-safe
        
    def register_handler(self, message_type: MessageType, handler: Callable):
        """Registra um handler para um tipo de mensagem"""
        self.message_handlers[message_type].append(handler)
    
    def connect_agent(self, agent: 'A2AAgent'):
        """Conecta outro agente para comunicação"""
        with self._lock:
            self.connected_agents[agent.agent_id] = agent
        
    def disconnect_agent(self, agent_id: str):
        """Desconecta um agente"""
        with self._lock:
            if agent_id in self.connected_agents:
                del self.connected_agents[agent_id]
    
    def send_message(self, message: A2AMessage) -> bool:
        """Envia mensagem para outro agente"""
        try:
            with self._lock:
                if message.receiver_id and message.receiver_id in self.connected_agents:
                    target_agent = self.connected_agents[message.receiver_id]
                    target_agent.receive_message(message)
                    self.message_history.append(message)
                    return True
                elif message.message_type == MessageType.BROADCAST:
                    # Broadcast para todos os agentes conectados
                    # Filtrar por tags se especificadas
                    target_tags = message.metadata.get('target_tags') if message.metadata else None
                    
                    for agent in self.connected_agents.values():
                        # Se target_tags especificadas, verificar se agente tem as tags
                        if target_tags:
                            # Simplificação: envia para todos (filtro seria implementado no agente)
                            agent.receive_message(message)
                        else:
                            agent.receive_message(message)
                    self.message_history.append(message)
                    return True
                return False
        except Exception as e:
            print(f"Erro ao enviar mensagem: {e}")
            return False
    
    def receive_message(self, message: A2AMessage):
        """Recebe e processa mensagem"""
        self.message_history.append(message)
        
        # Executa handlers registrados
        for handler in self.message_handlers[message.message_type]:
            try:
                handler(message)
            except Exception as e:
                print(f"Erro no handler: {e}")
    
    def create_request(self, receiver_id: str, action: str, params: Dict[str, Any]) -> A2AMessage:
        """Cria uma mensagem de requisição"""
        return A2AMessage.create(
            sender_id=self.agent_id,
            receiver_id=receiver_id,
            message_type=MessageType.REQUEST,
            content={
                "action": action,
                "params": params
            }
        )
    
    def create_response(self, original_message: A2AMessage, result: Any, success: bool = True) -> A2AMessage:
        """Cria uma mensagem de resposta"""
        return A2AMessage.create(
            sender_id=self.agent_id,
            receiver_id=original_message.sender_id,
            message_type=MessageType.RESPONSE,
            content={
                "result": result,
                "success": success
            },
            correlation_id=original_message.id
        )
    
    def broadcast(self, content: Dict[str, Any], target_tags: Optional[List[str]] = None) -> A2AMessage:
        """Cria uma mensagem de broadcast com filtro opcional por tags
        
        Args:
            content: Conteúdo da mensagem
            target_tags: Lista de tags para filtrar destinatários (opcional)
        """
        message = A2AMessage.create(
            sender_id=self.agent_id,
            message_type=MessageType.BROADCAST,
            content=content
        )
        
        # Adiciona tags ao metadata se especificadas
        if target_tags:
            if message.metadata is None:
                message.metadata = {}
            message.metadata['target_tags'] = target_tags
        
        self.send_message(message)
        return message

class A2AAgent:
    """Agente base com capacidades A2A"""
    
    def __init__(self, agent_id: str):
        self.agent_id = agent_id
        self.a2a_protocol = A2AProtocol(agent_id)
        self.setup_default_handlers()
    
    def setup_default_handlers(self):
        """Configura handlers padrão"""
        self.a2a_protocol.register_handler(MessageType.REQUEST, self.handle_request)
        self.a2a_protocol.register_handler(MessageType.RESPONSE, self.handle_response)
        self.a2a_protocol.register_handler(MessageType.NOTIFICATION, self.handle_notification)
    
    def handle_request(self, message: A2AMessage):
        """Handler padrão para requisições"""
        action = message.content.get("action")
        params = message.content.get("params", {})
        
        # Implementar lógica específica do agente aqui
        result = f"Ação '{action}' processada pelo agente {self.agent_id}"
        
        response = self.a2a_protocol.create_response(message, result)
        self.a2a_protocol.send_message(response)
    
    def handle_response(self, message: A2AMessage):
        """Handler padrão para respostas"""
        print(f"Resposta recebida de {message.sender_id}: {message.content}")
    
    def handle_notification(self, message: A2AMessage):
        """Handler padrão para notificações"""
        print(f"Notificação de {message.sender_id}: {message.content}")
    
    def receive_message(self, message: A2AMessage):
        """Recebe mensagem via protocolo A2A"""
        self.a2a_protocol.receive_message(message)
    
    def connect_to(self, other_agent: 'A2AAgent'):
        """Conecta-se a outro agente"""
        self.a2a_protocol.connect_agent(other_agent)
        other_agent.a2a_protocol.connect_agent(self)
    
    def send_request(self, receiver_id: str, action: str, params: Dict[str, Any]):
        """Envia requisição para outro agente"""
        message = self.a2a_protocol.create_request(receiver_id, action, params)
        return self.a2a_protocol.send_message(message)
    
    def notify_all(self, content: Dict[str, Any]):
        """Envia notificação para todos os agentes conectados"""
        return self.a2a_protocol.broadcast(content)