from sys import stdout
from time import sleep
from random import randint
from threading import Thread

from zmq import Poller, Context
from zmq.backend.cython import proxy
from zmq.backend.cython.constants import ROUTER, DEALER, POLLIN


def tprint(msg):
    stdout.write(msg + '\n')
    stdout.flush()


CLIENT_ADDR = 'eta'
CLIENT_PORT = '2960'
SERVER_ADDR = 'eta'
SERVER_PORT = '2961'


class Client(Thread):
    def __init__(self, key):
        self.id = key
        Thread.__init__ (self)

    def run(self):
        context = Context()
        socket = context.socket(DEALER)

        identity = f'worker-{self.id}'

        socket.identity = identity.encode('ascii')
        socket.connect(f'tcp://{CLIENT_ADDR}:{CLIENT_PORT}')

        print(f'Client {identity} started')
        poll = Poller()
        poll.register(socket, POLLIN)

        reqs = 0
        while True:
            reqs += 1
            print(f'Req {reqs} sent..')
            socket.send_string(f'request {reqs}')
            for i in range(5):
                sockets = dict(poll.poll(1000))
                if socket in sockets:
                    msg = socket.recv()
                    tprint(f'Client {identity} received: {msg}')

            if reqs == 5:
                break

        socket.close()
        context.term()


class Server(Thread):
    def __init__(self):
        Thread.__init__ (self)

    def run(self):
        context = Context()
        frontend = context.socket(ROUTER)
        frontend.bind(f'tcp://*:{CLIENT_PORT}')

        backend = context.socket(DEALER)
        backend.bind('inproc://backend')

        workers = []
        for i in range(5):
            worker = Worker(context)
            worker.start()
            workers.append(worker)

        proxy(frontend, backend)

        frontend.close()
        backend.close()
        context.term()


class Worker(Thread):
    def __init__(self, context):
        Thread.__init__ (self)
        self.context = context

    def run(self):
        worker = self.context.socket(DEALER)
        worker.connect('inproc://backend')
        tprint('Worker started')

        while True:
            ident, msg = worker.recv_multipart()
            tprint(f'Worker received {msg} from {ident}')
            replies = randint(0, 4)
            for i in range(replies):
                sleep(1. / (randint(1, 10)))
                worker.send_multipart([ident, msg])

        worker.close()


def main():
    server = Server()
    server.start()

    for i in range(3):
        client = Client(i)
        client.start()

    server.join()


if __name__ == "__main__":
    main()