"""
A2A (Agent-to-Agent) Protocol Manager.

Handles registration of local agents in a room and routes messages
between local and remote agents.
"""

import logging
import asyncio
import json
from typing import Dict, Any, List, Optional, Callable
from .a2a_models import AgentIdentity, A2AMessage

logger = logging.getLogger("brynq.rooms.a2a")

class A2AManager:
    def __init__(self, room_manager):
        self.room_manager = room_manager
        self.local_agents: Dict[str, AgentIdentity] = {}
        self.remote_agents: Dict[str, AgentIdentity] = {}
        self.callbacks: Dict[str, List[Callable]] = {} # agent_id -> list of handlers

    def register_local_agent(self, agent_id: str, name: str, capabilities: List[str], model: str, callback: Callable):
        """Register a local agent to participate in the room's A2A protocol."""
        prof = self.room_manager._get_local_profile()
        peer_id = prof.get("user_id", "unknown")
        
        identity = AgentIdentity(
            agent_id=agent_id,
            name=name,
            peer_id=peer_id,
            capabilities=capabilities,
            model=model
        )
        self.local_agents[agent_id] = identity
        
        if agent_id not in self.callbacks:
            self.callbacks[agent_id] = []
        self.callbacks[agent_id].append(callback)
        
        # Broadcast presence
        asyncio.create_task(self.broadcast_status(agent_id))
        logger.info(f"Local agent {name} ({agent_id}) registered for A2A")

    async def broadcast_status(self, agent_id: str):
        """Broadcast agent's status to the room."""
        if agent_id not in self.local_agents:
            return
        
        identity = self.local_agents[agent_id]
        msg = {
            "type": "agent_status",
            "agent_id": identity.agent_id,
            "name": identity.name,
            "peer_id": identity.peer_id,
            "capabilities": identity.capabilities,
            "model": identity.model,
            "status": identity.status
        }
        self.room_manager.send_message_raw(msg)

    def handle_incoming_message(self, data: Dict[str, Any]):
        """Process incoming A2A messages from the room WebSocket."""
        msg_type = data.get("type")
        
        if msg_type == "agent_status":
            self._handle_status(data)
        elif msg_type in ("agent_chat", "agent_task", "agent_result"):
            self._handle_message(data)

    def _handle_status(self, data: Dict[str, Any]):
        agent_id = data.get("agent_id")
        if not agent_id or agent_id in self.local_agents:
            return
            
        identity = AgentIdentity(
            agent_id=agent_id,
            name=data.get("name", "Unknown"),
            peer_id=data.get("peer_id", "unknown"),
            capabilities=data.get("capabilities", []),
            model=data.get("model", ""),
            status=data.get("status", "idle")
        )
        self.remote_agents[agent_id] = identity
        logger.debug(f"Discovered remote agent: {identity.name} ({agent_id})")

    def _handle_message(self, data: Dict[str, Any]):
        receiver_id = data.get("receiver_id")
        
        # If it's for everyone or specifically for a local agent
        if receiver_id == "all" or receiver_id in self.local_agents:
            # Deliver to local agent(s)
            target_ids = [receiver_id] if receiver_id != "all" else self.local_agents.keys()
            
            for tid in target_ids:
                if tid in self.callbacks:
                    for cb in self.callbacks[tid]:
                        try:
                            cb(data)
                        except Exception as e:
                            logger.error(f"Error in A2A callback for {tid}: {e}")

    def send_agent_message(self, sender_id: str, receiver_id: str, content: Any, msg_type: str = "agent_chat"):
        """Send a message from a local agent to another agent."""
        if sender_id not in self.local_agents:
            raise ValueError(f"Agent {sender_id} is not a registered local agent")
            
        msg = {
            "type": msg_type,
            "sender_id": sender_id,
            "receiver_id": receiver_id,
            "content": content,
            "timestamp": A2AMessage().timestamp,
            "id": A2AMessage().id
        }
        self.room_manager.send_message_raw(msg)
