import asyncio
import json

import zmq
from zmq.asyncio import Context

from control_backend.agents.base import BaseAgent
from control_backend.core.agent_system import InternalMessage
from control_backend.core.config import settings


class TestPauseAgent(BaseAgent):
    def __init__(self, name: str):
        super().__init__(name)
        
    async def setup(self):
        context = Context.instance()
        self.pub_socket = context.socket(zmq.PUB)
        self.pub_socket.connect(settings.zmq_settings.internal_pub_address)

        self.add_behavior(self._pause_command_loop())
        self.logger.debug("TestPauseAgent setup complete.")
    
    async def _pause_command_loop(self):
        print("Starting Pause command test loop.")
        while True:
            pause_command = {
                "endpoint": "pause",
                "data": True,
            }

            message = InternalMessage(
                to="ri_communication_agent",
                sender=self.name,
                body=json.dumps(pause_command),
            )
            await self.send(message)

            # User interrupt message 
            data = {
                "type": "pause",
                "context": True,
            }
            await self.pub_socket.send_multipart([b"button_pressed", json.dumps(data).encode()])

            self.logger.info("Pausing robot actions.")
            await asyncio.sleep(15)  # Simulate delay between messages

            pause_command = {
                "endpoint": "pause",
                "data": False,
            }
            message = InternalMessage(
                to="ri_communication_agent",
                sender=self.name,
                body=json.dumps(pause_command),
            )
            await self.send(message)

            # User interrupt message
            data = {
                "type": "pause",
                "context": False,
            }
            await self.pub_socket.send_multipart([b"button_pressed", json.dumps(data).encode()])

            self.logger.info("Resuming robot actions.")
            await asyncio.sleep(15)  # Simulate delay between messages
