"""
A2AAgent — A specialized BrynqAgent that participates in the room A2A protocol.
"""

import logging
import asyncio
from src.brynq.agent_loop import BrynqAgent
from runtime.rooms.room_manager import room_manager

logger = logging.getLogger("brynq.executor.a2a_agent")

class A2AAgent(BrynqAgent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.a2a_id = self.session_id # Use session_id as A2A identity
        
        # Register with the room A2A manager if we are in a room
        if room_manager.active_room:
            room_manager.a2a.register_local_agent(
                agent_id=self.a2a_id,
                name=self.name,
                capabilities=["chat", "code_review"], # Example capabilities
                model=self.model,
                callback=self._on_a2a_message
            )

    def _on_a2a_message(self, msg: dict):
        """Callback for incoming A2A messages."""
        msg_type = msg.get("type")
        sender = msg.get("sender_id")
        content = msg.get("content")
        
        logger.info(f"[{self.name}] Received A2A {msg_type} from {sender}: {content}")
        
        # Handle different types of messages
        if msg_type == "agent_task":
            # If we receive a task, process it and send back result
            asyncio.create_task(self._process_a2a_task(msg))

    async def _process_a2a_task(self, msg: dict):
        """Process an incoming task from another agent."""
        sender = msg.get("sender_id")
        task_id = msg.get("id")
        content = msg.get("content")
        
        logger.info(f"[{self.name}] Executing A2A task from {sender}...")
        
        # Call the real backend to solve the task
        try:
            response = self._call_backend(content)
            # Send result back
            room_manager.a2a.send_agent_message(
                sender_id=self.a2a_id,
                receiver_id=sender,
                content=response,
                msg_type="agent_result"
            )
        except Exception as e:
            logger.error(f"[{self.name}] A2A task execution failed: {e}")

    def send_to_agent(self, receiver_id: str, content: str):
        """Send a chat message to another agent in the room."""
        room_manager.a2a.send_agent_message(
            sender_id=self.a2a_id,
            receiver_id=receiver_id,
            content=content,
            msg_type="agent_chat"
        )

    def request_task(self, receiver_id: str, task_description: str):
        """Request another agent to perform a task."""
        room_manager.a2a.send_agent_message(
            sender_id=self.a2a_id,
            receiver_id=receiver_id,
            content=task_description,
            msg_type="agent_task"
        )
