import zmq

from robot_interface.endpoints.receiver_base import ReceiverBase
from robot_interface.state import state


class MainReceiver(ReceiverBase):
    def __init__(self, zmq_context, port=5555):
        """
        The main receiver endpoint, responsible for handling ping and negotiation requests.

        :param zmq_context: The ZeroMQ context to use.
        :type zmq_context: zmq.Context

        :param port: The port to use.
        :type port: int
        """
        super(MainReceiver, self).__init__("main", "json")
        self.create_socket(zmq_context, zmq.REP, port, bind=False)

    @staticmethod
    def _handle_ping(message):
        """A simple ping endpoint. Returns the provided data."""
        return {"endpoint": "ping", "data": message.get("data")}

    @staticmethod
    def _handle_port_negotiation(message):
        endpoints = [socket.endpoint_description() for socket in state.sockets]

        return {"endpoint": "negotiate/ports", "data": endpoints}

    @staticmethod
    def _handle_negotiation(message):
        """
        Handle a negotiation request. Will respond with ports that can be used to connect to the robot.

        :param message: The negotiation request message.
        :type message: dict

        :return: A response dictionary with a 'ports' key containing a list of ports and their function.
        :rtype: dict[str, list[dict]]
        """
        # In the future, the sender could send information like the robot's IP address, etc.

        if message["endpoint"] == "negotiate/ports":
            return MainReceiver._handle_port_negotiation(message)

        return {"endpoint": "negotiate/error", "data": "The requested endpoint is not implemented."}

    def handle_message(self, message):
        if "endpoint" not in message:
            return {"endpoint": "error", "data": "No endpoint provided."}

        if message["endpoint"] == "ping":
            return self._handle_ping(message)
        elif message["endpoint"].startswith("negotiation"):
            return self._handle_negotiation(message)

        return {"endpoint": "error", "data": "The requested endpoint is not supported."}