/*=========================================================================

  Program:   ParaView
  Module:    vtkZeroMQCommunicator.cxx

  Copyright (c) Kitware, Inc.
  All rights reserved.
  See Copyright.txt or http://www.paraview.org/HTML/Copyright.html for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/
#include "vtkZeroMQCommunicator.h"

#include "vtkCommand.h"
#include "vtkMultiProcessStream.h"
#include "vtkMutexLock.h"
#include "vtkObjectFactory.h"

#include <zmq.hpp>
#include <sstream>

void vtkSend(zmq::socket_t& socket, const void* data, size_t datasize)
{
  // So, to send a message, we:
  //  - Send an empty message frame with the MORE flag set; then
  //  - Send the message body.
  zmq::message_t empty(0);
  socket.send(empty, ZMQ_SNDMORE);

  zmq::message_t request(datasize);
  memcpy(request.data(), data, datasize);
//  cout << "Sent: " << datasize << endl;
  socket.send(request);
}

void vtkReceive(zmq::socket_t& socket, zmq::message_t &reply)
{
//  //And when we receive a message, we:
//  //  - Receive the first frame and if it's not empty, discard the whole message;
//  //  - Receive the next frame and pass that to the application.
  zmq::message_t empty;
  do
    {
    socket.recv(&empty);
    }
  while (empty.size() != 0);
  socket.recv(&reply);
///  cout << "Received: " << reply.size() << endl;
}

vtkStandardNewMacro(vtkZeroMQCommunicator);
//----------------------------------------------------------------------------
vtkZeroMQCommunicator::vtkZeroMQCommunicator()
  : Context(NULL),
  SocketSend(NULL),
  SocketReceive(NULL),
  Lock(vtkSimpleMutexLock::New())
{
}

//----------------------------------------------------------------------------
vtkZeroMQCommunicator::~vtkZeroMQCommunicator()
{
  delete this->SocketSend;
  this->SocketSend = NULL;
  delete this->SocketReceive;
  this->SocketReceive = NULL;
  this->Context = NULL;
  this->Lock->Delete();
  this->Lock = NULL;
}

//----------------------------------------------------------------------------
void vtkZeroMQCommunicator::SetContext(zmq::context_t *context)
{
  this->Lock->Lock();
  assert(context);
  this->Context = context;
  this->Lock->Unlock();
}

//----------------------------------------------------------------------------
bool vtkZeroMQCommunicator::Connect(const char* url)
{
  this->Lock->Lock();
  assert(this->SocketSend == NULL && this->SocketReceive == NULL && this->Context != NULL);

  // zeromq only support up to 256 chart in inproc url.
  assert(strlen(url) < 200);

  std::ostringstream ep_REP;
  ep_REP << url << "#0";
  this->SocketReceive = new zmq::socket_t(*this->Context, ZMQ_ROUTER);
  cout << "REP: " << ep_REP.str().c_str() << endl;
  this->SocketReceive->connect(ep_REP.str().c_str());

  std::ostringstream ep_REQ;
  ep_REQ << url << "#1";
  this->SocketSend = new zmq::socket_t(*this->Context, ZMQ_DEALER);
  cout << "REQ: " << ep_REQ.str().c_str() << endl;
  this->SocketSend->connect(ep_REQ.str().c_str());
  this->Lock->Unlock();
  return true;
}

//----------------------------------------------------------------------------
bool vtkZeroMQCommunicator::Bind(const char* url)
{
  this->Lock->Lock();
  assert(this->SocketSend == NULL && this->SocketReceive == NULL && this->Context != NULL);

  // zeromq only support up to 256 chart in inproc url.
  assert(strlen(url) < 200);

  std::ostringstream ep_REQ;
  ep_REQ << url << "#0";
  this->SocketSend = new zmq::socket_t(*this->Context, ZMQ_DEALER);
  cout << "REQ: " << ep_REQ.str().c_str() << endl;
  this->SocketSend->bind(ep_REQ.str().c_str());

  std::ostringstream ep_REP;
  ep_REP << url << "#1";
  this->SocketReceive = new zmq::socket_t(*this->Context, ZMQ_ROUTER);
  cout << "REP: " << ep_REP.str().c_str() << endl;
  this->SocketReceive->bind(ep_REP.str().c_str());
  this->Lock->Unlock();
  return true;
}

//----------------------------------------------------------------------------
int vtkZeroMQCommunicator::SendVoidArray(
  const void *data, vtkIdType length, int type, int remoteHandle, int tag)
{
  this->Lock->Lock();
  int typesize;
  switch (type)
    {
    vtkTemplateMacro(typesize = sizeof(VTK_TT));
  default:
    abort();
    }

  vtkMultiProcessStream stream;
  stream << length << type << tag;
  std::vector<unsigned char> sdata;
  stream.GetRawData(sdata);

  vtkSend(*this->SocketSend, &sdata[0], sdata.size());
  vtkSend(*this->SocketSend, data, length*typesize);
  this->Lock->Unlock();
  return 1;
}

//----------------------------------------------------------------------------
int vtkZeroMQCommunicator::ReceiveVoidArray(
  void *data, vtkIdType maxlength, int type, int remoteHandle, int tag)
{
  this->Lock->Lock();
  int typesize;
  switch (type)
    {
    vtkTemplateMacro(typesize = sizeof(VTK_TT));
  default:
    abort();
    }

  bool done = false;
  while (!done)
    {
    zmq::message_t metaMessage;
    zmq::message_t dataMesssage;

    vtkReceive(*this->SocketReceive, metaMessage);
    vtkReceive(*this->SocketReceive, dataMesssage);

    int ilength, itype, itag;
    vtkMultiProcessStream stream;
    stream.SetRawData(reinterpret_cast<const unsigned char*>(metaMessage.data()),
      static_cast<unsigned int>(metaMessage.size()));
    stream >> ilength >> itype >> itag;

    if (itag != tag)
      {
      // raise WrongTagEvent.
      char* idata = new char[dataMesssage.size() + sizeof(itag) + sizeof(int)];
      char* ptr = idata;
      memcpy(ptr, &itag, sizeof(itag));
      ptr += sizeof(itag);
      int length = dataMesssage.size();
      memcpy(ptr, &length, sizeof(length));
      ptr += sizeof(length);
      memcpy(ptr, dataMesssage.data(), dataMesssage.size());
      int res = this->InvokeEvent(vtkCommand::WrongTagEvent, idata);
      delete [] idata;
      if (res)
        {
        continue;
        }
      vtkErrorMacro("Tag mismatch: got " << itag << ", expecting " << tag << ".");
      return 0;
      }

    assert(tag == itag && type == itype && (maxlength*typesize) >= ilength &&
      dataMesssage.size() == ilength*typesize);
    memcpy(data, dataMesssage.data(), dataMesssage.size());

    this->Count = ilength;
    break;
    }
  this->Lock->Unlock();
  return 1;
}

//----------------------------------------------------------------------------
void vtkZeroMQCommunicator::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Lock->Lock();
  this->Superclass::PrintSelf(os, indent);
  this->Lock->Unlock();
}
