/****************************************************************************
**  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 <stdbool.h>
#include <stdint.h>

#include "sion_const.h"
#include "sion_debug.h"
#include "sion_error_handler.h"
#include "sion_fd.h"
#include "sion_file.h"
#include "sion_filedesc.h"
#include "sion_generic_apidesc.h"
#include "sion_generic_collective_merge.h"
#include "sion_internal.h"

/* callback functions */
int64_t _lsion_generic_collective_process_write_merge(const void *data, int64_t *spec, int sid);
bool _lsion_generic_collective_next_write_spec_merge(void *iterator_, int64_t *spec);

/* collective I/O */
#define DFUNCTION "_lsion_coll_fwrite_merge"
size_t _lsion_coll_fwrite_merge(const void *data, size_t size, size_t nitems, int sid)
{
  _sion_filedesc *sion_filedesc;
  _sion_generic_gendata *sion_gendata;
  _sion_generic_apidesc *sion_apidesc;

  if ((sid < 0) || (_lsion_vcdtype(sid) != SION_FILEDESCRIPTOR) || !(sion_filedesc = _lsion_vcdtovcon(sid))) {
    _lsion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_ABORT, "lsion_coll_fwrite: invalid sion_filedesc %d", sid);
    assert(0);
  }
  DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "enter usecoll=%d collector=%d collsize=%d (%d tasks, %d files)\n",
    sion_filedesc->usecoll, sion_filedesc->collector, sion_filedesc->collsize, sion_filedesc->ntasks, sion_filedesc->nfiles));

  sion_gendata = sion_filedesc->dataptr;
  sion_apidesc = sion_gendata->apidesc;

  /* needed for avoiding subsequent non-collective calls */
  sion_filedesc->collcmdused = 1;

  /* check collsize */
  if (sion_filedesc->collsize <= 0) {
    _lsion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_ABORT, "lsion_coll_fwrite: collsize=%d <= 0, returning ...\n",
      (int)sion_filedesc->collsize);
    assert(0);
  }

  /* check availability of gather_execute callback */
  if (!sion_apidesc->gather_execute_cb) {
    _lsion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_ABORT,
      "lsion_coll_fwrite: API %s not correctly initalized, collective I/O calls missing, aborting", sion_apidesc->name);
    assert(0);
  }

  /* parameter of callback function */
  int collector = (int)sion_filedesc->collector;
  int firstsender = collector + 1;
  int lastsender = collector + sion_filedesc->collsize - 1;
  if (lastsender > sion_filedesc->ntasks) {
    lastsender = sion_filedesc->ntasks - 1;
  }
  DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector=%d collsize=%d firstsender=%d lastsender=%d\n", collector,
    sion_filedesc->collsize, firstsender, lastsender));

  /* write own part */
  int64_t bytes_written;
  if (sion_filedesc->rank == sion_filedesc->collector) {
    DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "COLLECTOR write data of size %ld at current position\n", (long)nitems * size));
    bytes_written = _lsion_write_multi_chunk(sion_filedesc, data, nitems * size);

    /* collect and write parts of sender tasks via callback function */
    // TODO: handle return code
    sion_apidesc->gather_execute_cb(data, sion_filedesc->fsblksize, sion_gendata->comm_data_local, collector, firstsender,
      lastsender, sid, _lsion_generic_collective_process_write_merge, NULL, _lsion_generic_collective_next_write_spec_merge);

  } else {
    DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "SENDER send data of size %ld at current position\n", (long)nitems * size));

    int64_t iterator[] = {nitems * size, sion_filedesc->fsblksize};
    /* collect and write parts of sender tasks via callback function */
    bytes_written =
      sion_apidesc->gather_execute_cb(data, sion_filedesc->fsblksize, sion_gendata->comm_data_local, collector, firstsender,
        lastsender, sid, _lsion_generic_collective_process_write_merge, iterator, _lsion_generic_collective_next_write_spec_merge);
  }
  int64_t items_written = size ? bytes_written / size : 0;

  DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "leave usecoll=%d collector=%d collsize=%d (%d tasks, %d files) rc=%d\n",
    sion_filedesc->usecoll, sion_filedesc->collector, sion_filedesc->collsize, sion_filedesc->ntasks, sion_filedesc->nfiles,
    items_written));

  return items_written;
}
#undef DFUNCTION

#define DFUNCTION "_lsion_generic_collective_process_write_merge"
int64_t _lsion_generic_collective_process_write_merge(const void *data, int64_t *spec, int sid)
{
  _sion_filedesc *sion_filedesc;
  if ((sid < 0) || (_lsion_vcdtype(sid) != SION_FILEDESCRIPTOR) || !(sion_filedesc = _lsion_vcdtovcon(sid))) {
    return _lsion_errorprint(
      SION_SIZE_NOT_VALID, _SION_ERROR_RETURN, "_lsion_generic_collective_process_write_merge: invalid sion_filedesc %d", sid);
  }
  DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "enter spec[0]=%d spec[1]=%d sid=%d\n", (int)spec[0], (int)spec[1], sid));

  int64_t bytestowrite = spec[1];

  /* write data */
  DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector start to write data of size %lld at position %lld\n",
    (long long)bytestowrite, (long long)sion_filedesc->currentpos));
  int64_t bwrote = _lsion_write_multi_chunk(sion_filedesc, data, bytestowrite);
  if (bwrote != bytestowrite) {
    return _lsion_errorprint(
      bwrote, _SION_ERROR_RETURN, "_lsion_generic_collective_process_write_merge: problems writing data ...\n");
  }

  /* update internal data */
  DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector wrote data block (bwrote=%d) of size %ld new pos %lld, %lld\n", bwrote,
    (long)bytestowrite, (long long)sion_filedesc->currentpos, (long long)_lsion_file_get_position(sion_filedesc->fileptr)));

  DPRINTFP(
    (4, DFUNCTION, _SION_DEFAULT_RANK, "leave spec[0]=%d spec[1]=%d sid=%d rc=%d\n", (int)spec[0], (int)spec[1], sid, bwrote));
  return bwrote;
}
#undef DFUNCTION

bool _lsion_generic_collective_next_write_spec_merge(void *iterator_, int64_t *spec)
{
  int64_t *iterator = iterator_;
  if (iterator[0] < 1) {
    return false;
  }

  spec[0] = 0; // ignored by collector
  spec[1] = (iterator[0] > iterator[1]) ? iterator[1] : iterator[0];
  iterator[0] -= iterator[1];
  return true;
}
