#include <string>
#include <iostream>
#include <exception>
#include <thread>
#include <vector>
#include <zmq.hpp>
#include <stdio.h>
#include <assert.h>
#include <cstring>
#include <algorithm>
#include <cerrno>
#include <chrono>

#define SYNC_MSG "sync"
#define SYNC_MSGLEN strlen((SYNC_MSG))

using namespace std;

bool has_only_digits(const string s)
{
  return s.find_first_not_of( "0123456789" ) == string::npos;
}

// The function we want to execute on the new thread.
void task1(void * arg, int id)
{
  bool sync_done = false;

  zmq::context_t * context = (zmq::context_t *)arg;

  zmq::socket_t gettask_socket(*context, ZMQ_PULL);
  gettask_socket.connect("inproc://task_publisher");

  zmq::socket_t sendres_socket(*context, ZMQ_PUSH);
  sendres_socket.connect("inproc://collect_res");

  zmq::message_t sync_msg;
  gettask_socket.recv(&sync_msg); // block waiting for sync msg
  cout << "thread " << id << " receives: " << (char *)sync_msg.data() << endl;

  if(strlen((char *)sync_msg.data()) == SYNC_MSGLEN)
  {
    zmq::message_t msg_back((void *)to_string(id).c_str(), to_string(id).size() + 1, NULL);
    cout << "thread " << id << " sends: " << (char *)to_string(id).c_str() << endl;

    sendres_socket.send(msg_back);
    sync_done = true;
  }

  cout << "thread " << id << " sync done" << endl;

  // zmq::message_t msg(6);
  // memcpy((void *)msg.data(), "Hello", 6);
  // sendres_socket.send(msg);

  // zmq::message_t job;
  // gettask_socket.recv(&job);

  // cout << "in worker: " << (char *)job.data() << endl;

  // cout << "task1 says: Hello " << endl;
}

bool if_sync_done(int threadID, bool * sync_array, int worker_num)
{
  cout << "main thread receives sync msg from thread " << threadID << endl;

  // for(int i = 0; i < worker_num; i++)
  //   cout << "before " << boolalpha << sync_array[i] << " " << endl;

  if(!sync_array[threadID])
    sync_array[threadID] = true;

  // for(int i = 0; i < worker_num; i++)
  //   cout << "after " << boolalpha << sync_array[i] << " " << endl;

  bool sync_done = true;
  for(int i = 0; i < worker_num; i++)
    sync_done = sync_array[i] && sync_done;

  return sync_done;
}

int main(int argc, char* argv[])
{
  int worker_num;

  if(argc != 2)
  {
    cout << "1 parameter pls" << endl;
    exit(1);
  }

  try
  {
    if(!has_only_digits(string(argv[1])))
    {
      cout << "digit pls" << endl;
      exit(1);
    }
    worker_num = stoi(string(argv[1]));
  }
  catch(exception e)
  {
    cout << "exception while processing parameters" << endl;
    exit(1);
  }

  zmq::context_t context(0);
  zmq::socket_t distask_socket(context, ZMQ_PUSH);
  distask_socket.bind("inproc://task_publisher");

  zmq::socket_t getres_socket(context, ZMQ_PULL);
  getres_socket.bind("inproc://collect_res");

  vector<thread> pool;
  for(int i = 0; i < worker_num; i++)
  {
    cout << "main() : creating thread, " << i << endl;
    pool.push_back(thread(task1, (void *)&context, i));
  }

  bool sync_done = false;
  bool * sync_array = new bool[worker_num];
  fill_n(sync_array, worker_num, false);

  while(true)
  {
    if(sync_done)
    {
      cout << "sync done in main thread" << endl;
      break;
    }

    zmq::message_t sync_msg(4);
    memcpy((void *)sync_msg.data(), SYNC_MSG, SYNC_MSGLEN);
    for(int i = 0; i < worker_num; i++)
      distask_socket.send(sync_msg);

    for(int i = 0; i < worker_num; i++)
    {
      if(sync_done)
        break;
      if(i != 0)
        this_thread::sleep_for(chrono::milliseconds(500));

      zmq::message_t res_msg;
      int ret = getres_socket.recv(&res_msg, ZMQ_DONTWAIT);

      if(ret == -1 && errno == EAGAIN)
        continue;

      int threadID = stoi(string((char *)res_msg.data()));
      sync_done = if_sync_done(threadID, sync_array, worker_num);
    }
  }


  assert(sync_done);

  // for(int i = 0; i < worker_num; i++)
  // {
  //   zmq::message_t msg(6);
  //   memcpy((void *) msg.data (), "World", 6);
  //   distask_socket.send(msg);
  // }



  for(auto &t : pool)
    t.join();

  delete[] sync_array;

  exit(0);
}