// Copyright 2016 Intel Corporation
// SPDX-License-Identifier: Apache-2.0

#include "Context.h"
#include <snappy.h>
#include <algorithm>
#include <chrono>
#include <iostream>

#include "common/OSPCommon.h"
#include "rkcommon/memory/malloc.h"
#include "rkcommon/tasking/async.h"
#include "rkcommon/tasking/tasking_system_init.h"
#include "rkcommon/utility/getEnvVar.h"

using rkcommon::byte_t;
using rkcommon::make_unique;
using rkcommon::tasking::AsyncLoop;
using rkcommon::tasking::numTaskingThreads;
using rkcommon::utility::getEnvVar;

using namespace std::chrono;

namespace maml {

/*! the singleton object that handles all the communication */
std::unique_ptr<Context> Context::singleton;

Context::Context(bool enableCompression) : compressMessages(enableCompression)
{}

Context::~Context()
{
  try {
    stop();
  } catch (const std::exception &e) {
    ospray::handleError(OSP_UNKNOWN_ERROR, e.what());
  }
}

/*! register a new incoing-message handler. if any message comes in
  on the given communicator we'll call this handler */
void Context::registerHandlerFor(MPI_Comm comm, MessageHandler *handler)
{
  if (handlers.find(comm) != handlers.end()) {
    std::cerr << CODE_LOCATION
              << ": Warning: handler for this MPI_Comm already installed"
              << std::endl;
  }

  handlers[comm] = handler;

  /*! todo: to avoid race conditions we MAY want to check if there's
      any messages we've already received that would match this
      handler */
}

/*! put the given message in the outbox. note that this can be
  done even if the actual sending mechanism is currently
  stopped */
void Context::send(std::shared_ptr<Message> msg)
{
  // The message uses malloc/free, so use that instead of new/delete
  if (compressMessages) {
    byte_t *compressed =
        (byte_t *)malloc(snappy::MaxCompressedLength(msg->size));
    size_t compressedSize = 0;
    snappy::RawCompress(reinterpret_cast<const char *>(msg->data),
        msg->size,
        reinterpret_cast<char *>(compressed),
        &compressedSize);
    free(msg->data);

    msg->data = compressed;
    msg->size = compressedSize;
  }

  outbox.push_back(std::move(msg));
}

void Context::queueCollective(std::shared_ptr<Collective> col)
{
  // TODO WILL: auto-compress collectives?
  collectiveOutbox.push_back(std::move(col));
}

void Context::processInboxMessages()
{
  auto incomingMessages = inbox.consume();

  for (auto &msg : incomingMessages) {
    auto *handler = handlers[msg->comm];

    if (compressMessages) {
      // Decompress the message before handing it off
      size_t uncompressedSize = 0;
      snappy::GetUncompressedLength(reinterpret_cast<const char *>(msg->data),
          msg->size,
          &uncompressedSize);
      byte_t *uncompressed = (byte_t *)malloc(uncompressedSize);
      snappy::RawUncompress(reinterpret_cast<const char *>(msg->data),
          msg->size,
          reinterpret_cast<char *>(uncompressed));
      free(msg->data);

      msg->data = uncompressed;
      msg->size = uncompressedSize;
    }

    handler->incoming(msg);
  }
}

void Context::sendMessagesFromOutbox()
{
  auto outgoingMessages = outbox.consume();
  auto outgoingCollectives = collectiveOutbox.consume();

  for (auto &msg : outgoingMessages) {
    MPI_Request request;

    int rank = 0;
    MPI_CALL(Comm_rank(msg->comm, &rank));
    // Don't send to ourself, just forward to the inbox directly
    if (rank == msg->rank) {
      inbox.push_back(std::move(msg));
    } else {
      MPI_CALL(Isend(msg->data,
          msg->size,
          MPI_BYTE,
          msg->rank,
          msg->tag,
          msg->comm,
          &request));
      msg->started = high_resolution_clock::now();

      pendingSends.push_back(request);
      sendCache.push_back(std::move(msg));
    }
  }

  for (auto &col : outgoingCollectives) {
    col->start();
    pendingCollectives.push_back(col);
  }
}

void Context::pollForAndRecieveMessages()
{
  for (auto &it : handlers) {
    MPI_Comm comm = it.first;

    /* probe if there's something incoming on this handler's comm */
    int hasIncoming = 0;
    MPI_Status status;
    MPI_CALL(Iprobe(MPI_ANY_SOURCE, MPI_ANY_TAG, comm, &hasIncoming, &status));

    if (hasIncoming) {
      int size;
      MPI_CALL(Get_count(&status, MPI_BYTE, &size));

      auto msg = std::make_shared<Message>(size);
      msg->rank = status.MPI_SOURCE;
      msg->tag = status.MPI_TAG;
      msg->comm = comm;

      MPI_Request request;
      MPI_CALL(Irecv(
          msg->data, size, MPI_BYTE, msg->rank, msg->tag, msg->comm, &request));

      msg->started = high_resolution_clock::now();

      pendingRecvs.push_back(request);
      recvCache.push_back(std::move(msg));
    }
  }
}

void Context::waitOnSomeRequests()
{
  if (!pendingSends.empty() || !pendingRecvs.empty()) {
    const size_t totalMessages = pendingSends.size() + pendingRecvs.size();
    int *done = STACK_BUFFER(int, totalMessages);
    MPI_Request *mergedRequests = STACK_BUFFER(MPI_Request, totalMessages);

    for (size_t i = 0; i < totalMessages; ++i) {
      if (i < pendingSends.size()) {
        mergedRequests[i] = pendingSends[i];
      } else {
        mergedRequests[i] = pendingRecvs[i - pendingSends.size()];
      }
    }

    int numDone = 0;
    MPI_CALL(Testsome(
        totalMessages, mergedRequests, &numDone, done, MPI_STATUSES_IGNORE));

    for (int i = 0; i < numDone; ++i) {
      size_t msgId = done[i];
      if (msgId < pendingSends.size()) {
        pendingSends[msgId] = MPI_REQUEST_NULL;
        sendCache[msgId] = nullptr;
      } else {
        msgId -= pendingSends.size();

        inbox.push_back(std::move(recvCache[msgId]));

        pendingRecvs[msgId] = MPI_REQUEST_NULL;
        recvCache[msgId] = nullptr;
      }
    }

    // Clean up anything we sent
    sendCache.erase(std::remove(sendCache.begin(), sendCache.end(), nullptr),
        sendCache.end());

    pendingSends.erase(
        std::remove(pendingSends.begin(), pendingSends.end(), MPI_REQUEST_NULL),
        pendingSends.end());

    // Clean up anything we received
    recvCache.erase(std::remove(recvCache.begin(), recvCache.end(), nullptr),
        recvCache.end());

    pendingRecvs.erase(
        std::remove(pendingRecvs.begin(), pendingRecvs.end(), MPI_REQUEST_NULL),
        pendingRecvs.end());
  }
  if (!pendingCollectives.empty()) {
    pendingCollectives.erase(std::remove_if(pendingCollectives.begin(),
                                 pendingCollectives.end(),
                                 [](const std::shared_ptr<Collective> &col) {
                                   return col->finished();
                                 }),
        pendingCollectives.end());
  }
}

void Context::flushRemainingMessages()
{
  while (!pendingRecvs.empty() || !pendingSends.empty()
      || !pendingCollectives.empty() || !inbox.empty() || !outbox.empty()) {
    sendMessagesFromOutbox();
    pollForAndRecieveMessages();
    waitOnSomeRequests();
    processInboxMessages();
  }
}

/*! start the service; from this point on maml is free to use MPI
  calls to send/receive messages; if your MPI library is not
  thread safe the app should _not_ do any MPI calls until 'stop()'
  has been called */
void Context::start()
{
  std::lock_guard<std::mutex> lock(tasksMutex);
  if (!isRunning()) {
    tasksAreRunning = true;

    auto launchMethod = AsyncLoop::LaunchMethod::AUTO;

    auto MAML_SPAWN_THREADS = getEnvVar<int>("MAML_SPAWN_THREADS");
    if (MAML_SPAWN_THREADS) {
      launchMethod = MAML_SPAWN_THREADS.value()
          ? AsyncLoop::LaunchMethod::THREAD
          : AsyncLoop::LaunchMethod::TASK;
    }

    if (!sendReceiveThread.get()) {
      sendReceiveThread = make_unique<AsyncLoop>(
          [&]() {
            sendMessagesFromOutbox();
            pollForAndRecieveMessages();
            waitOnSomeRequests();
          },
          launchMethod);
    }

    if (!processInboxThread.get()) {
      processInboxThread = make_unique<AsyncLoop>(
          [&]() { processInboxMessages(); }, launchMethod);
    }

    sendReceiveThread->start();
    processInboxThread->start();
  }
}

bool Context::isRunning() const
{
  return tasksAreRunning;
}

/*! stops the maml layer; maml will no longer perform any MPI calls;
  if the mpi layer is not thread safe the app is then free to use
  MPI calls of its own, but it should not expect that this node
  receives any more messages (until the next 'start()' call) even
  if they are already in fligh
  WILL: Don't actually stop, for reasons described above flush messages
*/
void Context::stop()
{
  std::lock_guard<std::mutex> lock(tasksMutex);
  if (tasksAreRunning) {
    quitThreads = true;
    if (sendReceiveThread) {
      sendReceiveThread->stop();
    }
    if (processInboxThread) {
      processInboxThread->stop();
    }

    tasksAreRunning = false;
    flushRemainingMessages();
  }
}

} // namespace maml