from typing import List
import zmq

from mininet.common.node import Node
from mininet.models.models import Address


class Broker(Node):
    def __init__(self):
        super().__init__()
        self.poller = zmq.Poller()

    def init_socket(self, addr: Address | List[Address]):
        self.socket = [self.context.socket(zmq.ROUTER), self.context.socket(zmq.DEALER)]
        self.socket[0].bind(addr[0].format())
        self.socket[1].bind(addr[1].format())
        self.poller.register(self.socket[0], zmq.POLLIN)
        self.poller.register(self.socket[1], zmq.POLLIN)

    def run(self):
        while True:
            if self._stop:
                return
            else:
                socks = dict(self.poller.poll())
                if socks.get(self.socket[0]) == zmq.POLLIN:
                    message = self.socket[0].recv_multipart()
                    self.socket[1].send_multipart(message)
                if socks.get(self.socket[1]) == zmq.POLLIN:
                    message = self.socket[1].recv_multipart()
                    self.socket[0].send_multipart(message)

class Boss(Node):
    def __init__(self):
        super().__init__()
        self.reply: str | None = None

    def init_socket(self, addr: Address | List[Address]):
        self.socket = self.context.socket(zmq.REQ)
        self.socket.connect(addr.format())

    def run(self):
        while True:
            if self._stop:
                return
            else:
                self.socket.send_string("Hello")
                self.reply = self.socket.recv()

class Worker(Node):
    def __init__(self):
        super().__init__()
        self.request: str | None = None

    def init_socket(self, addr: Address | List[Address]):
        self.socket = self.context.socket(zmq.REP)
        self.socket.bind(addr.format())

    def run(self):
        while True:
            if self._stop:
                return
            else:
                self.request = self.socket.recv()
                self.socket.send_string("World")
