#!/usr/bin/env perl
# wesyoung.me / csirtgadgets.org
# updated to use LibZMQ4 - Brian Kelly / fullauto.com
use 5.011; # so we can use 'say'
use strict;
use warnings;
use ZMQ::LibZMQ4;
use ZMQ::Constants qw(:all);
use threads;
use Carp::Assert;
use constant NBR_CLIENTS    => 7;
use constant NBR_WORKERS    => 2;
use constant READY          => "\001";
use constant BACKEND_URL   => 'ipc://backend.ipc';
use constant FRONTEND_URL  => 'ipc://frontend.ipc';
main();
sub client_thread {
    my $i   = shift;
    my $id = 'Client-'.$i;
    my $ctx = zmq_init();
    my $socket = zmq_socket($ctx, ZMQ_REQ);
    my $rv = zmq_setsockopt($socket, ZMQ_IDENTITY, $id);
    assert($rv == 0, 'setting socket options');
    $rv    = zmq_connect($socket,FRONTEND_URL());
    assert($rv == 0, 'connecting client…');
    my $reply;
    while(1){
        say "$id sending Hello";
        my $rv = zmq_msg_send('Hello',$socket);
        assert($rv == 5, 'sending hello');
        say "$id waiting for reply…";
        $reply = zmq_recvmsg($socket);
        assert($reply);
        say "$id got a reply";
    }
}
sub worker_thread {
    my $i   = shift;
    my $id = 'Worker-'.$i;
    my $ctx = zmq_init();
    my $socket = zmq_socket($ctx, ZMQ_REQ);
    my $rv = zmq_setsockopt($socket, ZMQ_IDENTITY, $id);
    assert($rv == 0);
    $rv    = zmq_connect($socket,BACKEND_URL());
    assert($rv == 0,'connecting to backend');
    say "$id sending READY";
    $rv = zmq_msg_send(READY(), $socket);
    assert($rv);
    my $buf = zmq_msg_init();
    while(1){
       say "$id waiting…";
       my @msg=();
       while (1) {
          my $rv = zmq_msg_recv($buf,$socket);
          push @msg, zmq_msg_data($buf) if $rv;
          while (zmq_getsockopt($socket,ZMQ_RCVMORE)) {
             $rv = zmq_msg_recv($buf,$socket);
             push @msg, zmq_msg_data($buf);
          }
          last if $rv==-1 || -1<$#msg;
       }
       if (-1<$#msg) {
          say "$id got: $msg[2] from $msg[0]";
          say "$id sending OK to $msg[0]";
          zmq_msg_send($msg[0],$socket,ZMQ_SNDMORE);
          zmq_msg_send('',$socket,ZMQ_SNDMORE);
          zmq_msg_send('OK',$socket);
       }
    }
}
sub main {
    say 'starting main…';
    my $client_nbr  = NBR_CLIENTS();
    my $ctx = zmq_init();
    my $frontend = zmq_socket($ctx, ZMQ_ROUTER);
    my $backend  = zmq_socket($ctx, ZMQ_ROUTER);
    my $rv  = zmq_bind($frontend,FRONTEND_URL());
    assert($rv == 0);
    $rv     = zmq_bind($backend,BACKEND_URL());
    assert($rv == 0);
    my @threads;
    foreach my $t (1 .. NBR_WORKERS()){
        say 'starting worker: '.$t;
        push(@threads,threads->create('worker_thread',$t));
    }
    foreach my $t (1 .. NBR_CLIENTS()){
        say 'starting client: '.$t;
        push(@threads,threads->create('client_thread',$t));
    }
    my @workers;
    my ($w_addr,$delim,$c_addr,$data);
    my $items = [
        {
            events      => ZMQ_POLLIN,
            socket      => $frontend,
            callback    => sub {
                if (-1<$#workers) {
                    say 'frontend…';
                    my $buf = zmq_msg_init();
                    my @msg=();
                    while (1) {
                       my $rv = zmq_msg_recv($buf,$frontend);
                       push @msg, zmq_msg_data($buf) if $rv;
                       while (zmq_getsockopt($frontend,ZMQ_RCVMORE)) {
                          $rv = zmq_msg_recv($buf,$frontend);
                          push @msg, zmq_msg_data($buf);
                       }
                       last if $rv==-1 || -1<$#msg;
                    }
                    assert($#msg);
                    zmq_msg_send(pop(@workers),$backend,ZMQ_SNDMORE);
                    zmq_msg_send('',$backend,ZMQ_SNDMORE);
                    zmq_msg_send($msg[0],$backend,ZMQ_SNDMORE);
                    zmq_msg_send('',$backend,ZMQ_SNDMORE);
                    zmq_msg_send($msg[2],$backend,);
                }
            },
        },
        {
            events      => ZMQ_POLLIN,
            socket      => $backend,
            callback    => sub {
                say 'backend…';
                my $buf = zmq_msg_init();
                my @msg=();
                while (1) {
                   my $rv = zmq_msg_recv($buf,$backend);
                   push @msg, zmq_msg_data($buf) if $rv;
                   while (zmq_getsockopt($backend,ZMQ_RCVMORE)) {
                      $rv = zmq_msg_recv($buf,$backend);
                      push @msg, zmq_msg_data($buf);
                   }
                   last if $rv==-1 || -1<$#msg;
                }
                assert($#msg);
                assert($#workers < NBR_WORKERS());
                $w_addr = $msg[0];
                push(@workers,$w_addr);
                $delim = $msg[1];
                assert($delim eq '');
                $c_addr = $msg[2];
                if($c_addr ne READY()){
                    $delim = $msg[3];
                    assert ($delim eq '');
                    $data = $msg[4];
                    say 'sending '.$data.' to '.$c_addr;
                    zmq_msg_send($c_addr,$frontend,ZMQ_SNDMORE);
                    zmq_msg_send('',$frontend,ZMQ_SNDMORE);
                    zmq_msg_send($data,$frontend);
                } else {
                    say 'worker checking in: '.$w_addr;
                }
            },
        },
    ];
    while(1){ zmq_poll($items); select undef,undef,undef,0.025; }
    $_->join() for(@threads);
}