#
#   Request-reply service in Python
#   Connects REP socket to tcp://localhost:5560
#   Expects "Hello" from client, replies with "World"
#
import zmq

# context = zmq.Context)
context = zmq.Context.instance()

frontend = context.socket(zmq.ROUTER)
frontend.bind("tcp://127.0.0.1:5559")
backend = context.socket(zmq.DEALER)
backend.bind("tcp://127.0.0.1:5560")
# Initialize poll set
poller = zmq.Poller()
poller.register(frontend, zmq.POLLIN)
poller.register(backend, zmq.POLLIN)

while True:
    socks = dict(poller.poll())

    if socks.get(frontend) == zmq.POLLIN:
        # why frontend.recv() not work! For the envelop reason
        message = frontend.recv_multipart()
        # [b'\x00k\x8bEg', b'', b'Hello 2']
        print('what msg in ROUTER recv? ', message)
        backend.send_multipart(message)

    if socks.get(backend) == zmq.POLLIN:
        message = backend.recv_multipart()
        print('what msg in DEALER recv? ', message)
        frontend.send_multipart(message)

import zmq
import threading
import random
import time

context = zmq.Context()
DEALER_NUM = 5
sinks = []
for _ in range(DEALER_NUM):
    sinks.append(context.socket(zmq.DEALER))

for i in range(DEALER_NUM):
    sinks[i].connect("inproc://example")

def pull_worker(*args):
    idx = args[0]
    total = 0
    start = time.time()
    while True:
        msg = sinks[idx].recv_multipart()
        time.sleep( 1 * random.randint(1, 3))
        total += 1
        if total == 9:
            end = time.time()
            print('elapse time(%s) seconds for idx(%d) queue' % (end-start, idx))

threads = []
for _ in range(DEALER_NUM):
    threads.append( threading.Thread(target=pull_worker, args=[_]) )
for i in range(DEALER_NUM):
    threads[i].start()

anonymous = context.socket(zmq.PUSH)
anonymous.bind("inproc://example")

for i in range(DEALER_NUM * 10):
    anonymous.send(b"gogogo")

while True:
    time.sleep(1)

#
# N REQ ->1 ROUTER to  demonstrate loading balance strategy
#
import zmq
import threading
import random
import time

context = zmq.Context()

PULL_NUM = 5

sinks = []
for _ in range(PULL_NUM):
    sinks.append(context.socket(zmq.REQ))

for i in range(PULL_NUM):
    sinks[i].connect("inproc://example")


def worker(*args):
    idx = args[0]
    total = 0
    start = time.time()
    while True:
        sinks[idx].send(b'ready')

        msg = sinks[idx].recv()
        if msg == b'end':
            end = time.time()
            print('elapse time(%s) seconds for idx(%d) queue, processed %d tasks' % (end - start, idx,total))
            break

        time.sleep(1 * random.randint(1, 3))
        # time.sleep(0.1)
        total += 1

threads = []
for _ in range(PULL_NUM):
    threads.append(threading.Thread(target=worker, args=[_]))
for i in range(PULL_NUM):
    threads[i].start()

anonymous = context.socket(zmq.ROUTER)
anonymous.bind("inproc://example")

for i in range(PULL_NUM * 10):
    [addr, empty, msg] = anonymous.recv_multipart()
    # print('addr %s msg %s' % (addr, msg))
    if msg == b'ready':
        # ROUTE send in `3 frames` format
        anonymous.send_multipart([addr, empty, b'gogogo'])

for i in range(PULL_NUM):
    [addr, empty, msg] = anonymous.recv_multipart()
    anonymous.send_multipart([addr, empty, b'end'])

# LRU queue using pyzmq IOLoop reactor
import threading
import time
from zmq.eventloop.ioloop import IOLoop
from zmq.eventloop.zmqstream import ZMQStream
import zmq
import random

NBR_CLIENTS = 50
NBR_WORKERS = 5

context = zmq.Context.instance()
localfd_url = "tcp://127.0.0.1:6000"
localbd_url = "tcp://127.0.0.1:6001"


def worker_task(ident):
    """
    REQ as worker
    """
    worker = context.socket(zmq.REQ)
    worker.identity = u"Worker-{}".format(ident).encode("ascii")
    worker.connect(localbd_url)

    # Tell broker , I'm ready
    worker.send(b'READY')

    count = 0

    total = 0

    # Processing tasks
    while True:
        try:
            start = time.time()
            [addr, empty, message] = worker.recv_multipart()
            time.sleep(random.randint(0, 3) * 0.1)
            # time.sleep(random.randint(0, 3) * 0.05)
            reply = u"Worker-{} process done".format(ident).encode("ascii")
            worker.send_multipart([addr, empty, reply])
            count += 1
            end = time.time()
            total += (end -start)
        except zmq.ContextTerminated:
            print('Worker-%d got %s tasks , spent %f seconds' % (ident, count, total))
            return

    pass


def client_task(ident):
    """
    REQ as client
    :param id:
    :return:
    """
    client = context.socket(zmq.REQ)
    client.connect(localfd_url)
    client.identity = u"Client-{}".format(ident).encode("ascii")
    # Send req and expect a rep , an normal REQ-REP flow
    client.send(b'Hello there, i am client-%d ' % ident)
    try:
        reply = client.recv()
        print('client-%d got reply:%s' % (ident, reply))
    except zmq.ContextTerminated:
        return


def single_broker_instance():
    localfd = context.socket(zmq.ROUTER)
    localbd = context.socket(zmq.ROUTER)

    localfd.bind(localfd_url)
    localbd.bind(localbd_url)

    # Add to poll
    poller = zmq.Poller()
    # Only poll for requests from backend until workers are available
    poller.register(localbd, zmq.POLLIN)

    workers = []
    work_ready = 0

    count = 0

    worked_started = 0

    while True:
        sockets = dict(poller.poll())

        # worker may got READY or worker reply
        if localbd in sockets:
            requests = localbd.recv_multipart()
            address, empty, messages = requests[:3]
            if not workers:
                poller.register(localfd, zmq.POLLIN)
            # Worker ready
            workers.append(address)
            if messages == b'READY' and len(requests) == 3:
                print('%s ready' % address)

                # Start clients then
                if worked_started == NBR_WORKERS - 1:
                    for i in range(NBR_CLIENTS):
                        threading.Thread(target=client_task, args=[i]).start()

                worked_started += 1

            # Messages sent from(ROUTER) to client be  [client address, empty, sth] , good for ROUTER routing.
            elif len(requests) == 5:
                client_address = messages
                empty, messages = requests[3:]
                localfd.send_multipart([client_address, empty, messages])
                count += 1
            else:
                print('some error')
                break
            # print(count)
            if count == NBR_CLIENTS - 1:
                print('quit')
                break

        # Got req from clients ROUTER->ROUTER
        if localfd in sockets:
            [address, empty, messages] = localfd.recv_multipart()
            print('valid workers:', workers)
            worker_address = workers.pop(0)
            localbd.send_multipart([worker_address, empty, address, empty, messages])
            if not workers:
                poller.unregister(localfd)
            pass

    localbd.close()
    localfd.close()
    context.term()
    pass

def broker_tasks(ident):
    pass


if __name__ == '__main__':


    for i in range(NBR_WORKERS):
        threading.Thread(target=worker_task, args=[i]).start()

    single_broker_instance()