import json
import logging
from spade.agent import Agent
from spade.behaviour import CyclicBehaviour
import zmq

from control_backend.core.config import settings
from control_backend.core.zmq_context import context
from control_backend.schemas.message import Message

logger = logging.getLogger(__name__)

class TestAgent(Agent):
    socket: zmq.Socket

    class ListenBehaviour(CyclicBehaviour):
        async def run(self):
            assert self.agent is not None
            topic, body = await self.agent.socket.recv_multipart()
            
            try:
                message_json = json.loads(body.decode("utf-8"))
                message = Message.model_validate(message_json)
                logger.info("Received message \"%s\"", message.message)
            except Exception as e:
                logger.error("Error processing message: %s", e)

    async def setup(self):
        logger.info("Setting up %s", self.jid)
        self.socket = context.socket(zmq.SUB)
        self.socket.connect(settings.zmq_settings.internal_comm_address)
        self.socket.setsockopt(zmq.SUBSCRIBE, b"message")

        b = self.ListenBehaviour()
        self.add_behaviour(b)

        logger.info("Finished setting up %s", self.jid)