//
// mineview collaboration server
// Each client connects to the receiver socket/
//   - Sends "Hello", and receives a unique ID from the server.
//   - Sends position/orientation updates, and receives "OK" from the server.
// Each client connects to the publisher socket, and receives
//   position/orientation updates from all connected clients, including itself.
//
// The basic message format is one of the two following formats
//   HelloPVMZ -- returns a unique client ID
//   PMVZ <clientID> <sessionID> <msgType> <message payload>

#include "zhelpers.hpp"

#ifndef NDEBUG
int64_t last_time = 0;
#endif

bool validID(
  zmq::socket_t &receiver,
  std::string &client,
  unsigned int maxClientID
  )
{
  client = s_recv(receiver);
  unsigned int clientID = maxClientID + 1;
  bool valid = false;
  try
  {
    unsigned int clientID = std::stoi(client);
    valid = clientID <= maxClientID;
    if (!valid)
    {
      throw nullptr;
    }
  }
  catch (...)
  {
#ifndef NDEBUG
    std::cout << "Invalid client ID received, ignoring message " << std::endl;
    s_dump(receiver);
#else
    s_clear(receiver);
#endif
    s_send(receiver, "ERROR");
  }
  return valid;
}

void processGenericMessage(
  zmq::socket_t &publisher,
  zmq::socket_t &receiver,
  unsigned int maxClientID)
{
  std::string client;
  bool publishReply = validID(receiver, client, maxClientID);
  if (!publishReply) {
    return;
  }
    // only retrieve the session if we got a client ID.
  std::string session = s_recv(receiver);

  // publish to all clients, use client-supplied session header
  // so clients can subscribe to different sessions.
  s_sendmore(publisher, session);
  // send client ID.
  s_sendmore(publisher, client);

#ifndef NDEBUG
  // Use timer, don't flood the output
  int64_t curr_time = s_clock();
  if (curr_time - last_time > 1000)
  {
    last_time = curr_time;
    std::cout << "Valid message from " << client
      << " for session " << session << std::endl;
  }
#endif

  int more;
  size_t more_size = sizeof(more);
  // a well-formatted message will always have another part, but guard against
  // bad-actors.
  receiver.getsockopt(ZMQ_RCVMORE, &more, &more_size);
  while (more)
  {
    zmq::message_t message;

    //  Process all parts of the message
    receiver.recv(&message);
    receiver.getsockopt(ZMQ_RCVMORE, &more, &more_size);

    publisher.send(message, more? ZMQ_SNDMORE: 0);
    // if (!more) done, last message part
  }

  //  Send reply back to client
  // s_send(receiver, "OK");
}

int main (int argc, char *argv[])
{
  zmq::context_t context(1);

  //  Connect to client - ID and updates from client.
  zmq::socket_t receiver(context, ZMQ_ROUTER);
  receiver.bind("tcp://*:5555");

  //  publish socket for broadcasting updates.
  zmq::socket_t publisher(context, ZMQ_PUB);
  publisher.bind("tcp://*:5556");

  //  Initialize poll set
  zmq::pollitem_t items [] =
  {
    { receiver, 0, ZMQ_POLLIN, 0 },
    { publisher, 0, ZMQ_POLLIN, 0 }
  };

  unsigned int currID = 1;
  //  Process messages from both sockets
  while (1)
  {
    zmq::poll(&items[0], 2, -1);

    if (items[0].revents & ZMQ_POLLIN)
    {
      // router socket must copy identity, so it goes to the right client (dealer)
      zmq::message_t identity;
      if(!receiver.recv(&identity, ZMQ_DONTWAIT))
      {
        // we thought there was a message, but nothing is ready.
        continue;
      }
      std::string msg = s_recv(receiver);

      if (msg == "HelloPMVZ")
      {
        std::stringstream ss;
        ss << currID;
#ifndef NDEBUG
        // send back the ID the client should use.
        std::cout << "Received: " << msg << " Sending ID: " << ss.str() << std::endl;
#endif
        receiver.send(identity, ZMQ_SNDMORE);
        s_send(receiver, ss.str());
        currID += 1;
      }
      else if (msg == "ping")
      {
        // ping-pong heartbeat - client sends pings when not sending any other data,
        // so can determine if server is still alive.
        std::string client;
        bool publishReply = validID(receiver, client, currID-1);
        if (publishReply) {
          receiver.send(identity, ZMQ_SNDMORE);
          s_send(receiver, "pong");
        }
      }
      else
      {
        // follow https://github.com/booksbyus/zguide/blob/master/examples/C%2B%2B/wuproxy.cpp
        // we are going to check for a valid header - a client ID we have sent
        // then forward the rest of the message pieces to our publish socket.
        if (msg != "PMVZ")
        {
#ifndef NDEBUG
          std::cout << "Missing header, ignoring message " << std::endl;
          s_dump(receiver);
#else
          s_clear(receiver);
#endif
        }
        else
        {
          processGenericMessage(publisher, receiver, currID-1);
        }
      }
    }
    if (items [1].revents & ZMQ_POLLIN)
    {
      // unneeded - publish socket can't recv.
      // publisher.recv(&message);
      //  Process update
    }
    s_sleep(0);
  }
  return 0;
}
