/****************************************************************************
**  SIONLIB     http://www.fz-juelich.de/jsc/sionlib                       **
*****************************************************************************
**  Copyright (c) 2008-2019                                                **
**  Forschungszentrum Juelich, Juelich Supercomputing Centre               **
**                                                                         **
**  See the file COPYRIGHT in the package base directory for details       **
****************************************************************************/

#define _XOPEN_SOURCE 700

#include <assert.h>
#include <inttypes.h>
#include <mpi.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

#define USE_PMPIno
#ifdef USE_PMPI
#define MPI_Comm_rank PMPI_Comm_rank
#define MPI_Comm_size PMPI_Comm_size
#define MPI_Gather PMPI_Gather
#define MPI_Scatter PMPI_Scatter
#define MPI_Bcast PMPI_Bcast
#define MPI_Barrier PMPI_Barrier
#define MPI_Comm_split PMPI_Comm_split
#define MPI_Send PMPI_Send
#define MPI_Recv PMPI_Recv
#endif

#include "sion_const.h"
#include "sion_debug.h"
#include "sion_error_handler.h"
#include "sion_filedesc.h"
#include "sion_ompi_cb_gen.h"
#include "sion_ompi_internal_gen.h"

#ifdef SION_OMPI

static void *__ompicol_global_pointer;

void *__sion_ompicol_share_ptr(void *in_ptr);

#define DFUNCTION "_ompi_gather_execute_cb"
int64_t _sion_ompi_gather_process_cb(const void *indata, int64_t fsblksize, void *commdata, int collector, int range_start,
  int range_end, int sid, int64_t process_cb(const void *, int64_t *, int), void *iterator,
  bool iterator_next_cb(void *, int64_t *))
{
  _ompi_api_commdata *sapi = (_ompi_api_commdata *)commdata;
  int rank = sapi->rank;

  DPRINTFP(
    (256, DFUNCTION, rank, " input collector=%d range_start=%d range_end=%d sid=%d\n", collector, range_start, range_end, sid));

  /* STEP1: collect info on thread level: iterators -> pointer to iterator on thread, indatas -> pointer to indata on thread */
  void **iterators = NULL;
#pragma omp master
  iterators = malloc(sapi->num_threads * sizeof(void *));

  iterators = __sion_ompicol_share_ptr(iterators);

  if (iterators == NULL) {
    _sion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_ABORT,
      "_sion_ompi_gather_execute_cb: cannot allocate temporary memory of size %zu (helpdata), aborting ...\n",
      (size_t)sapi->num_threads * sizeof(void *));
    assert(0);
  }

  /* gather iterators */
  iterators[sapi->thread_num] = iterator;

  void const **indatas = NULL;
#pragma omp master
  indatas = malloc(sapi->num_threads * sizeof(const void *));

  indatas = __sion_ompicol_share_ptr(indatas);

  if (indatas == NULL) {
    free(iterators);
    _sion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_ABORT,
      "_sion_ompi_gather_execute_cb: cannot allocate temporary memory of size %zu (tcounts), aborting ...\n",
      (size_t)sapi->num_threads * sizeof(const void *));
    assert(0);
  }

  /* gather pointers to data buffers */
  indatas[sapi->thread_num] = indata;

  int64_t *bytes_writtens = NULL;
#pragma omp master
  bytes_writtens = malloc(sapi->num_threads * sizeof(int64_t));

  bytes_writtens = __sion_ompicol_share_ptr(bytes_writtens);

  if (bytes_writtens == NULL) {
    free(iterators);
    free(indatas);
    _sion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_ABORT,
      "_sion_ompi_gather_execute_cb: cannot allocate temporary memory of size %zu (bytes_writtens), aborting ...\n",
      (size_t)sapi->num_threads * sizeof(int64_t *));
    assert(0);
  }

// make sure iterators and indatas have been gathered
#pragma omp barrier

/* STEP2: proceed on MPI level, master threads have all needed info */
#pragma omp master
  {
    MPI_Comm commgroup = sapi->comm;

    if (rank == collector) {
      /* it's the collector */
      int mrank = _sion_map_rank_ompi_to_mpi(rank, sapi->num_threads);

      /* allocate buffer */
      char *buffer = malloc(fsblksize);
      if (buffer == NULL) {
        _sion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_ABORT,
          "_sion_ompi_gather_process_cb: cannot allocate temporary memory of size %" PRId64 " (buffer), aborting ...\n",
          fsblksize);
        assert(0);
      }

      /* scan all other tasks */
      for (int t = range_start; t <= range_end; ++t) {
        int mt = _sion_map_rank_ompi_to_mpi(t, sapi->num_threads);

        if (mt == mrank) {
          /* thread is on same MPI rank */
          int tt = _sion_map_rank_ompi_to_thread_num(t, sapi->num_threads);

          /* process data directly */
          const char *p = indatas[tt];
          void *titerator = iterators[tt];
          int64_t spec[] = {-1, -1};
          int64_t total_bytes_written = 0;
          while (iterator_next_cb(titerator, spec)) {
            int64_t bytes_written = process_cb(p, spec, sid);
            p += bytes_written;
            total_bytes_written += bytes_written;
            if (bytes_written != spec[1]) {
              break;
            }
          }
          bytes_writtens[tt] = total_bytes_written;
        } else {
          /* loop over data parts */
          while (true) {
            /* receive spec */
            int64_t spec[] = {-1, -1};
            DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector wait for spec from %d\n", mt));
            MPI_Recv(spec, 2, MPI_INT64_T, mt, 1534, commgroup, MPI_STATUS_IGNORE);
            DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector got spec from %d (%lld,%lld)\n", mt, (long long)spec[0],
              (long long)spec[1]));

            if (spec[0] == -1) {
              DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector got end of stream\n"));
              break;
            }

            /* receive portion or all data */
            DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector wait for data block from %d\n", mt));
            assert(spec[1] <= fsblksize);
            MPI_Recv(buffer, spec[1], MPI_CHAR, mt, 1536, commgroup, MPI_STATUS_IGNORE);
            DPRINTFP(
              (4, DFUNCTION, _SION_DEFAULT_RANK, "collector got data block from %d datasize=%lld\n", mt, (long long)spec[1]));

            /* process data with callback function */
            DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector process data block of size %lld at pos %lld\n",
              (long long)spec[1], (long long)spec[0]));

            int64_t bytes_written = process_cb(buffer, spec, sid);

            // report back result of processing block
            DPRINTFP(
              (4, DFUNCTION, _SION_DEFAULT_RANK, "collector reporting result to %d result=%lld\n", mt, (long long)bytes_written));
            MPI_Send(&bytes_written, 1, MPI_INT64_T, mt, 1535, commgroup);
          }
        } /* not on local MPI rank */
      }

      free(buffer);
    } else {
      /* it's a sender */

      int mcollector = _sion_map_rank_ompi_to_mpi(collector, sapi->num_threads);

      /* send data for all threads on MPI rank */
      for (int tt = 0; tt < sapi->num_threads; ++tt) {
        int64_t spec[] = {-1, -1};
        const char *p = indatas[tt];
        void *titerator = iterators[tt];
        int64_t total_bytes_written = 0;
        while (iterator_next_cb(titerator, spec)) {
          /* send spec to collector */
          DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "sender send spec to %d (%lld,%lld)\n", collector, (long long)spec[0],
            (long long)spec[1]));
          MPI_Send(spec, 2, MPI_INT64_T, mcollector, 1534, commgroup);

          /* send data in chunks of fsblksize */
          DPRINTFP(
            (4, DFUNCTION, _SION_DEFAULT_RANK, "sender send data block to %d of size %lld\n", mcollector, (long long)spec[1]));
          MPI_Send(p, spec[1], MPI_CHAR, mcollector, 1536, commgroup);

          int64_t bytes_written = 0;
          MPI_Recv(&bytes_written, 1, MPI_INT64_T, mcollector, 1535, commgroup, MPI_STATUS_IGNORE);
          p += bytes_written;
          total_bytes_written += bytes_written;
          if (bytes_written != spec[1]) {
            DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "sender got short write count from %d write_count=%lld\n", mcollector,
              (long long)bytes_written));
            break;
          }
        }

        // signal end of stream to collector
        spec[0] = -1;
        spec[1] = -1;
        DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "sender send end of stream to %d (%lld,%lld)\n", mcollector,
          (long long)spec[0], (long long)spec[1]));
        MPI_Send(spec, 2, MPI_INT64_T, mcollector, 1534, commgroup);
        bytes_writtens[tt] = total_bytes_written;
      } /* for all threads */
    } /* sender */
  } /* omp master */

// make sure the master thread has updated all write counts
#pragma omp barrier

  int64_t rc = (sapi->thread_num == 0) ? 0 : bytes_writtens[sapi->thread_num];

// make sure everyone has read their write count
#pragma omp barrier

#pragma omp master
  {
    free(indatas);
    free(iterators);
    free(bytes_writtens);
  }

  DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "leave collector=%d rc=%d\n", collector, rc));

  return rc;
}
#undef DFUNCTION

#define DFUNCTION "_ompi_process_scatter_cb"
int64_t _sion_ompi_process_scatter_cb(void *outdata, int64_t fsblksize, void *commdata, int collector, int range_start,
  int range_end, int sid, int64_t process_cb(void *, int64_t *, int), void *iterator, bool iterator_next_cb(void *, int64_t *))
{
  _ompi_api_commdata *sapi = (_ompi_api_commdata *)commdata;
  int rank = sapi->rank;

  DPRINTFP(
    (256, DFUNCTION, rank, " input collector=%d range_start=%d range_end=%d sid=%d\n", collector, range_start, range_end, sid));

  /* STEP1: collect info on thread level: iterators -> pointer to iterator on thread, outdatas -> pointer to outdata on thread */
  void **iterators = NULL;
#pragma omp master
  iterators = malloc(sapi->num_threads * sizeof(void *));

  iterators = __sion_ompicol_share_ptr(iterators);

  if (iterators == NULL) {
    _sion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_ABORT,
      "_sion_ompi_process_scatter_cb: cannot allocate temporary memory of size %zu (iterators), aborting ...\n",
      (size_t)sapi->num_threads * sizeof(void *));
    assert(0);
  }

  /* gather iterators */
  iterators[sapi->thread_num] = iterator;

  void **outdatas = NULL;
#pragma omp master
  outdatas = malloc(sapi->num_threads * sizeof(void *));

  outdatas = __sion_ompicol_share_ptr(outdatas);

  if (outdatas == NULL) {
    free(iterators);
    _sion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_ABORT,
      "_sion_ompi_process_scatter_cb: cannot allocate temporary memory of size %zu (outdatas), aborting ...\n",
      (size_t)sapi->num_threads * sizeof(void *));
    assert(0);
  }

  /* store info about spec */
  outdatas[sapi->thread_num] = outdata;

  int64_t *bytes_reads = NULL;
#pragma omp master
  bytes_reads = malloc(sapi->num_threads * sizeof(int64_t *));

  bytes_reads = __sion_ompicol_share_ptr(bytes_reads);

  /* check return code from malloc */
  if (bytes_reads == NULL) {
    free(iterators);
    free(outdatas);
    _sion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_ABORT,
      "_sion_ompi_process_scatter_cb: cannot allocate temporary memory of size %zu (bytes_reads), aborting ...\n",
      (size_t)sapi->num_threads * sizeof(int64_t *));
    assert(0);
  }

// make sure iterators and outdatas have been gathered
#pragma omp barrier

/* STEP2: proceed on MPI level, master threads have all needed info  */
#pragma omp master
  {
    MPI_Comm commgroup = sapi->comm;

    if (rank == collector) {
      /* it's the collector */
      /* allocate buffer */
      char *buffer = malloc(fsblksize);
      if (buffer == NULL) {
        _sion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_ABORT,
          "_ompi_gather_process_cb: cannot allocate temporary memory of size %" PRId64 " (buffer), aborting ...\n", fsblksize);
        assert(0);
      }

      int mrank = _sion_map_rank_ompi_to_mpi(rank, sapi->num_threads);

      /* scan all other tasks */
      for (int t = range_start; t <= range_end; ++t) {
        int mt = _sion_map_rank_ompi_to_mpi(t, sapi->num_threads);

        if (mt == mrank) {
          /* thread is on same MPI rank */
          int tt = _sion_map_rank_ompi_to_thread_num(t, sapi->num_threads);
          /* process data directly */
          char *p = outdatas[tt];
          void *titerator = iterators[tt];
          int64_t spec[] = {-1, -1};
          int64_t total_bytes_read = 0;
          while (iterator_next_cb(titerator, spec)) {
            int64_t bytes_read = process_cb(p, spec, sid);
            p += bytes_read;
            total_bytes_read += bytes_read;
            if (bytes_read != spec[1]) {
              break;
            }
          }
          bytes_reads[tt] = total_bytes_read;
        } else {
          /* loop over data parts */
          while (true) {
            /* receive spec */
            int64_t spec[] = {-1, -1};
            DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector wait for spec from %d\n", t));
            MPI_Recv(spec, 2, MPI_INT64_T, mt, 1534, commgroup, MPI_STATUS_IGNORE);
            DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector got spec from %d (%lld,%lld)\n", t, (long long)spec[0],
              (long long)spec[1]));

            if (spec[0] == -1) {
              DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector got end of stream\n"));
              break;
            }

            /* process data with callback function */
            DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector process data block of size %lld at pos %lld\n",
              (long long)spec[1], (long long)spec[0]));
            int64_t bytes_read = process_cb(buffer, spec, sid);

            /* send portion or all data */
            DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector send for data block to %d\n", mt));
            MPI_Send(buffer, bytes_read, MPI_CHAR, mt, 1536, commgroup);
            DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector sent data block to %d datasize=%lld bytestorecv=%lld\n", mt,
              (long long)bytes_read));
          } /* while */
        } /* not on local MPI rank */
      } /* for all tasks */

      free(buffer);
    } else {
      /* it's a sender */

      int mcollector = _sion_map_rank_ompi_to_mpi(collector, sapi->num_threads);

      /* send data for all threads on MPI rank */
      for (int tt = 0; tt < sapi->num_threads; ++tt) {
        int64_t spec[] = {-1, -1};
        char *p = outdatas[tt];
        void *titerator = iterators[tt];
        int64_t total_bytes_read = 0;
        while (iterator_next_cb(titerator, spec)) {
          /* send spec to collector */
          DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "sender send spec to %d (%lld,%lld)\n", mcollector, (long long)spec[0],
            (long long)spec[1]));
          MPI_Send(spec, 2, MPI_INT64_T, mcollector, 1534, commgroup);

          /* send data in chunks of fsblksize */
          DPRINTFP(
            (4, DFUNCTION, _SION_DEFAULT_RANK, "sender recv data block from %d of size %lld\n", mcollector, (long long)spec[1]));
          MPI_Status status;
          MPI_Recv(p, spec[1], MPI_CHAR, mcollector, 1536, commgroup, &status);
          int bytes_read;
          MPI_Get_count(&status, MPI_CHAR, &bytes_read);
          p += bytes_read;
          total_bytes_read += bytes_read;

          if (bytes_read != spec[1]) {
            DPRINTFP(
              (4, DFUNCTION, _SION_DEFAULT_RANK, "sender got short read count from %d read_count=%d\n", mcollector, bytes_read));
            break;
          }
        } /* while */

        spec[0] = -1;
        spec[1] = -1;
        DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "sender send end of stream to %d (%lld,%lld)\n", mcollector,
          (long long)spec[0], (long long)spec[1]));
        MPI_Send(spec, 2, MPI_INT64_T, mcollector, 1534, commgroup);
        bytes_reads[tt] = total_bytes_read;
      } /* for all threads */
    } /* sender */
  } /* omp master */

// make sure the master thread has updated all read counts
#pragma omp barrier

  int64_t rc = (sapi->thread_num == 0) ? 0 : bytes_reads[sapi->thread_num];

// make sure everyone has read their write count
#pragma omp barrier

#pragma omp master
  {
    free(outdatas);
    free(iterators);
    free(bytes_reads);
  }

  DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "leave collector=%d rc=%d\n", collector, rc));

  return rc;
}
#undef DFUNCTION

/* share in_ptr given on master with all other threads, return value is the shared ptr */
#define DFUNCTION "__sion_ompi_share_ptr"
void *__sion_ompicol_share_ptr(void *in_ptr)
{
// make sure everyone is done with the global pointer
#pragma omp barrier

#pragma omp master
  __ompicol_global_pointer = in_ptr;

// make sure the master thread has updated the global pointer
#pragma omp barrier

  void *out_ptr = __ompicol_global_pointer;

// make sure everyone has read the global pointer
#pragma omp barrier
  return out_ptr;
}
#undef DFUNCTION

#define DFUNCTION "_sion_ompi_get_capability_cb"
int _sion_ompi_get_capability_cb(void *commdata)
{
  int rc = SION_CAPABILITY_NONE;
  _ompi_api_commdata *sapi = (_ompi_api_commdata *)commdata;

  if (sapi->thread_num == 0) {
    rc = SION_CAPABILITY_FULL;
    DPRINTFP((256, DFUNCTION, sapi->rank, "FULL capability\n"));
  } else {
    rc = SION_CAPABILITY_ONLY_SENDER;
    DPRINTFP((256, DFUNCTION, sapi->rank, "ONLY SENDER capability\n"));
  }
  return rc;
}
#undef DFUNCTION
/* end of ifdef SION_OMPI */
#endif
