import random

import time

import threading

import zmq

stop = False

READY = "READY"

context = zmq.Context()

def collector(frontend, backend):

  poller = zmq.Poller()

  poller.register(backend, zmq.POLLIN)

  backends = set()

  while not stop:

    frontend.send_multipart(["",READY])

    poller.register(frontend, zmq.POLLIN)

    for socket, event in poller.poll(100):

      request = socket.recv_multipart()

      if socket is backend:

        if request[2] == READY:

          backends.add(request[0])

      else:

        timeout = int(request[request.index("")+1])

        recipients = backends

        backends = set()

        for dest in recipients:

          backend.send_multipart([dest] + request)

        poller.unregister(frontend)

        try:

          while not stop and recipients:

            for socket, event in poller.poll(timeout):

              reply = socket.recv_multipart()

              if reply[2] == READY:

                backends.add(reply[0])

                recipients.discard(reply[0])

              else:

                frontend.send_multipart(reply[2:])

        except ZMQError:

          pass

def broker_collector(frontend_url, backend_url):

  frontend = context.socket(zmq.XREP)

  backend = context.socket(zmq.XREP)

  frontend.bind(frontend_url)

  backend.bind(backend_url)

  collector(frontend, backend)

def proxy_collector(frontend_url, backend_url):

  frontend = context.socket(zmq.XREP)

  backend = context.socket(zmq.XREP)

  frontend.connect(frontend_url)

  backend.bind(backend_url)

  collector(frontend, backend)

def worker(socket, workload, failure_rate = 0):

  while not stop:

    socket.send_multipart(["",READY])

    request = socket.recv_multipart()

    delim = request.index("")

    timeout = request[delim+1]

    request = request[delim+2:]

    assert request[0] == "REQUEST"

    if failure_rate and random.randrange(failure_rate) == 0:

      return

    time.sleep(workload)

    socket.send_multipart(request[:delim+1] + ["DONE"])

def connect_worker(url, workload, failure_rate = 0):

  while not stop:

    socket = context.socket(zmq.XREQ)

    socket.connect(url)

    worker(socket, workload, failure_rate)

def requester(socket, timeout = -1):

  while not stop:

    socket.send_multipart(["", str(timeout), "REQUEST"])

    results = 0

    while True:

      reply = socket.recv_multipart()

      if reply == ["",READY]:

        break

      results += 1

    print results, "results received"

  socket.send_multipart(["", "STOP"])

def connect_requester(url, timeout):

  socket = context.socket(zmq.XREQ)

  socket.connect(url)

  requester(socket, timeout)

if __name__ == "__main__":

  feurl = "inproc://frontend"

  beurl = "inproc://backend"

  brokers = []

  broker = threading.Thread(target = broker_collector, args = (feurl, beurl))

  broker.start()

  brokers.append(broker)

  time.sleep(1)

  senders = []

  for sender in xrange(2):

    sender = threading.Thread(target = connect_requester, args = (feurl,1000))

    sender.start()

    senders.append(sender)

  proxies = []

  proxy_urls = []

  for proxy in xrange(2):

    url = "inproc://proxy_be#%d" % (proxy,)

    proxy = threading.Thread(target = proxy_collector, args = (beurl, url))

    proxy.start()

    proxies.append(proxy)

    proxy_urls.append(url)

  time.sleep(1)

  workers = []

  for url in proxy_urls:

    for work in xrange(5):

      work = threading.Thread(target = connect_worker, args = (url, 0.1, 100))

      work.start()

      workers.append(work)

  time.sleep(10)

  stop = True

  for thread in senders + brokers + proxies + workers:

    thread.join()