from abc import abstractmethod
import zmq
import sys
from zmq.eventloop.zmqstream import ZMQStream
import threading
from mininet.models.models import Address
from tornado.ioloop import IOLoop


class Node:
    def __init__(self):
        self._stopping = threading.Event()
        self._stream: ZMQStream | None = None
        self._thread: threading.Thread | None = None
        self.context = zmq.Context.instance()
        self.loop = IOLoop.instance()

    @abstractmethod
    def init_socket(self):
        """Initialize all sockets needed by the subclass"""

    @abstractmethod
    def start(self):
        """
        Start node doing whatever it does, potentially blocking
        :return:
        """

    def start_threaded(self):
        self._thread = threading.Thread(target=self.start, daemon=True)
        self._thread.start()



    def send(self, msg: str):
        raise NotImplementedError('Class does not implement send')

    def term(self):
        self.context.term()

    def stop(self):
        self._stopping.set()


    def receive(self, msg: list[bytes | zmq.Frame]):
        """Do something to handle a incoming message"""
        raise NotImplementedError()

    def _wrap_socket(self, socket: zmq.Socket) -> ZMQStream:
        """Wrap a zmq socket and set its callback method"""
        stream = ZMQStream(socket, self.loop)
        stream.on_recv(self.receive)
        return stream


class Proxy(Node):
    def __init__(self, addr_fe: Address, addr_be: Address):
        self.addr_fe = addr_fe
        self.addr_be = addr_be
        super().__init__()


    def init_socket(self):
        self.frontend = self.context.socket(zmq.XSUB)
        self.frontend.connect(self.addr_fe.format())
        self.backend = self.context.socket(zmq.XPUB)
        self.backend.bind(self.addr_be.format_bind())

        self.frontend.setsockopt(zmq.SUBSCRIBE, b'A')

    def forward(self):
        zmq.proxy(self.frontend, self.backend)

    def start(self):
        self.init_socket()
        self.forward()

class Worker(Node):
    def __init__(self, addr: Address):
        self.addr = addr
        super().__init__()

    def init_socket(self):
        self.publisher = self.context.socket(zmq.PUB)
        self.publisher.bind(self.addr.format_bind())

    def send_heartbeat(self):
        while True:
            self.publisher.send_string("Jetson Orin Nano")

    # def start(self):
    #     self.send_heartbeat()

    def start(self):
        self._stopping.clear()

        self.init_socket()
        # self._stream = self._wrap_socket(self.publisher)

    def send(self, msg: bytes):
        self.publisher.send(msg)

    def exit(self):
        self.publisher.close()
        super().term()

class Client(Node):
    def __init__(self, addr: Address):
        self.addr = addr
        self.worker_list = list()
        super().__init__()

    def init_socket(self):
        self.subscriber = self.context.socket(zmq.SUB)
        self.subscriber.connect(self.addr.format())
        self.subscriber.setsockopt(zmq.SUBSCRIBE, b'A')


    def receive(self, msg: list[bytes | zmq.Frame]):
        # in the future when we have multiple kinds of messages, something like
        # if msg.type == "etc":
        #    self.handle_etc()
        # but actually write it like
        # self.handlers[msg.type](msg)

        self.worker_list.append(msg)


    # def recv_heartbeat(self):
    #     while True:
    #         if self.subscriber:
    #             self.worker_list.append(self.subscriber.recv_string())

    def start(self):
        self._stopping.clear()

        self.init_socket()
        while not self._stopping.is_set():
            print('AAAAAA')
            sys.stdout.flush()
            msg = self.subscriber.recv()
            self.receive(msg)
        #
        # self._stream = self._wrap_socket(self.subscriber)
        #
        # try:
        #     self.loop.start()
        # except RuntimeError:
        #     pass


    def exit(self):
        self.subscriber.close()
        super().term()
