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

#define _XOPEN_SOURCE 700

#include <assert.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_mpi_cb_gen.h"

#ifdef SION_MPI

#define DFUNCTION "_mpi_gather_execute_cb"
int64_t _lsion_mpi_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 *))
{
  int size, rank;
  _mpi_api_commdata *sapi = (_mpi_api_commdata *)commdata;
  MPI_Comm commp = sapi->comm;

  MPI_Comm_rank(commp, &rank);
  MPI_Comm_size(commp, &size);

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

  if (rank == collector) {
    /* it's the collector */

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

    /* scan all other tasks */
    for (int t = range_start; t <= range_end; ++t) {
      while (true) {
        // TODO: change these into a pair of Irecv (block) then Recv (spec)
        int64_t spec[] = {-1, -1};

        /* receive spec */
        DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector wait for spec from %d\n", t));
        MPI_Recv(spec, 2, MPI_INT64_T, t, 1534, commp, 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;
        }

        /* receive portion or all data */
        DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector wait for data block from %d\n", t));
        MPI_Recv(buffer, spec[1], MPI_CHAR, t, 1536, commp, MPI_STATUS_IGNORE);
        DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector got data block from %d datasize=%lld\n", t, (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", t, (long long)bytes_written));
        MPI_Send(&bytes_written, 1, MPI_INT64_T, t, 1535, commp);
      }
    }
    /* remove buffer */
    free(buffer);
    return 0;
  } else if ((rank >= range_start) && (rank <= range_end)) {
    /* it's a sender */
    int64_t spec[] = {-1, -1};
    const char *p = indata;
    int64_t total_bytes_written = 0;

    while (iterator_next_cb(iterator, spec)) {
      // TODO: change these two into a pair of Ssend (spec) then Rsend (data)
      /* 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, collector, 1534, commp);

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

      int64_t bytes_written = 0;
      MPI_Recv(&bytes_written, 1, MPI_INT64_T, collector, 1535, commp, 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", collector,
          (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", collector, (long long)spec[0],
      (long long)spec[1]));
    MPI_Send(spec, 2, MPI_INT64_T, collector, 1534, commp);

    DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "leave collector=%d rc=%d\n", collector, total_bytes_written));
    return total_bytes_written;
  } else {
    DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "task take not part in collective operation\n"));
    return 0;
  }
}
#undef DFUNCTION

#define DFUNCTION "_mpi_process_scatter_cb"
int64_t _lsion_mpi_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 *))
{
  int size, rank;
  _mpi_api_commdata *sapi = (_mpi_api_commdata *)commdata;
  MPI_Comm commp = sapi->comm;

  MPI_Comm_rank(commp, &rank);
  MPI_Comm_size(commp, &size);

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

  if (rank == collector) {
    /* it's the collector */

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

    /* scan all other tasks */
    for (int t = range_start; t <= range_end; ++t) {
      while (true) {
        int64_t spec[] = {-1, -1};

        /* receive spec */
        DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector wait for spec from %d\n", t));
        MPI_Recv(spec, 2, MPI_INT64_T, t, 1534, commp, 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]));

        /* send signal to send data */
        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", t));
        MPI_Send(buffer, bytes_read, MPI_CHAR, t, 1536, commp);
        DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector sent data block to %d datasize=%lld\n", t, (long long)bytes_read));
      }
    }
    /* remove buffer */
    free(buffer);
    return 0;
  } else if ((rank >= range_start) && (rank <= range_end)) {
    /* it's a sender */
    int64_t spec[] = {-1, -1};
    char *p = outdata;
    int64_t total_bytes_read = 0;

    while (iterator_next_cb(iterator, 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, collector, 1534, commp);

      /* receive data in chunks of fsblksize */
      DPRINTFP(
        (4, DFUNCTION, _SION_DEFAULT_RANK, "sender recv data block from %d of size %lld\n", collector, (long long)spec[1]));
      MPI_Status status;
      MPI_Recv(p, spec[1], MPI_CHAR, collector, 1536, commp, &status);
      int bytes_read;
      MPI_Get_count(&status, MPI_CHAR, &bytes_read);
      DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "sender recv data block from %d of size %lld (%d)\n", collector,
        (long long)spec[1], 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", collector, bytes_read));
        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", collector, (long long)spec[0],
      (long long)spec[1]));
    MPI_Send(spec, 2, MPI_INT64_T, collector, 1534, commp);

    DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "leave collector=%d rc=%d\n", collector, total_bytes_read));
    return total_bytes_read;
  } else {
    DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "task take not part in collective operation\n"));
    return 0;
  }
}
#undef DFUNCTION

#define DFUNCTION "_lsion_ompi_get_capability_cb"
int _lsion_mpi_get_capability_cb(void *commdata)
{
  int rc = SION_CAPABILITY_NONE;
  ONLY_DEBUG(_mpi_api_commdata *sapi = (_mpi_api_commdata *)commdata;)

  rc = SION_CAPABILITY_FULL;
  DPRINTFP((256, DFUNCTION, sapi->rank, "FULL capability\n"));
  return rc;
}
#undef DFUNCTION

/* end of ifdef MPI */
#endif
