/****************************************************************************
**  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 <inttypes.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

#include "sion_common.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_buddy.h"
#include "sion_generic_collective.h"
#include "sion_internal.h"
#include "sion_internal_positioning.h"

/* collective I/O */
#define DFUNCTION "lsion_coll_fwrite"
size_t lsion_coll_fwrite(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;

  /* no collective mode */
  if (!sion_filedesc->usecoll) {
    return lsion_fwrite(data, size, nitems, sid);
  }

  /* branch to merge mode if enabled */
  if (sion_filedesc->collmergemode) {
    return _lsion_coll_fwrite_merge(data, size, nitems, sid);
  }

  /* 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));

  int64_t bytes_written;
  if (sion_filedesc->rank == sion_filedesc->collector) {
    /* write own part */
    bytes_written = _lsion_write_multi_chunk(sion_filedesc, data, nitems * size);
    int64_t ownnewposition = sion_filedesc->currentpos;

    // 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, NULL, _lsion_generic_collective_next_write_spec);

    /* set own position to end of own block written in this call */
    _lsion_file_set_position(sion_filedesc->fileptr, ownnewposition);
    sion_filedesc->currentpos = ownnewposition;
  } else {
    _sion_partitioned_range_iterator ranges =
      _lsion_partitioned_range_iterator_from_filedesc(sion_filedesc, size * nitems, sion_filedesc->fsblksize);

    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, &ranges, _lsion_generic_collective_next_write_spec);

    /* set file pointer in data structure and in file if it is exported and can be used without control of SIONlib */
    _lsion_write_dry_run(sion_filedesc, nitems * size);
  }
  int64_t items_written = size ? bytes_written / size : 0;

  /* switch to buddy checkpointing, if enabled */
  if (sion_filedesc->usebuddy) {
    // TODO: how to include buddy result in return value
    _lsion_coll_fwrite_buddy(data, size, nitems, sid, sion_gendata);
  }

  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_coll_fread"
size_t lsion_coll_fread(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_fread: 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;

  /* no collective mode */
  if (!sion_filedesc->usecoll) {
    return lsion_fread(data, size, nitems, sid);
  }

  /* switch to buddy checkpointing, if enabled */
  if (sion_filedesc->usebuddy) {
    return _lsion_coll_fread_buddy(data, size, nitems, sid);
  }

  /* 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_fread: collsize=%d <= 0, returning ...\n", (int)sion_filedesc->collsize);
    assert(0);
  }

  // check availability of execute_scatter callback
  if (!sion_apidesc->execute_scatter_cb) {
    _lsion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_ABORT,
      "lsion_coll_fread: 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 = sion_filedesc->rank + 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));

  int64_t bytes_read;
  if (sion_filedesc->rank == sion_filedesc->collector) {
    /* read own part */
    bytes_read = _lsion_read_multi_chunk(sion_filedesc, data, nitems * size);
    int64_t ownnewposition = sion_filedesc->currentpos;

    /* read parts and scatter these to sender tasks via callback function */
    // TODO: handle return code
    sion_apidesc->execute_scatter_cb(data, sion_filedesc->fsblksize, sion_gendata->comm_data_local, collector, firstsender,
      lastsender, sid, _lsion_generic_collective_process_read, NULL, _lsion_generic_collective_next_read_spec);

    /* set own position to end of own block read in this call */
    _lsion_file_set_position(sion_filedesc->fileptr, ownnewposition);
    sion_filedesc->currentpos = ownnewposition;
  } else {
    _sion_partitioned_limited_range_iterator ranges =
      _lsion_partitioned_limited_range_iterator_from_filedesc(sion_filedesc, size * nitems, sion_filedesc->fsblksize);

    /* read parts and scatter these to sender tasks via callback function */
    bytes_read = sion_apidesc->execute_scatter_cb(data, sion_filedesc->fsblksize, sion_gendata->comm_data_local, collector,
      firstsender, lastsender, sid, _lsion_generic_collective_process_read, &ranges, _lsion_generic_collective_next_read_spec);
    bytes_read = nitems * size;

    /* set file pointer in data structure and in file if it is exported and can be used without control of SIONlib */
    _lsion_read_dry_run(sion_filedesc, nitems * size);
  }
  int64_t items_read = size ? bytes_read / 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_read));

  return items_read;
}
#undef DFUNCTION

#define DFUNCTION "_lsion_generic_collective_process_write"
int64_t _lsion_generic_collective_process_write(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: 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));

  /* move file pointer */
  int64_t destpos = spec[0], bytestowrite = spec[1];
  DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "currentpos=%d destpos=%d\n", (int)sion_filedesc->currentpos, (int)destpos));
  if (sion_filedesc->currentpos != destpos) {
    _lsion_file_set_position(sion_filedesc->fileptr, destpos);
    sion_filedesc->currentpos = destpos;
  }

  /* 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)destpos));
  int64_t bwrote = _lsion_file_write(data, bytestowrite, sion_filedesc->fileptr);
  if (bwrote != bytestowrite) {
    return _lsion_errorprint(bwrote, _SION_ERROR_RETURN, "_lsion_generic_collective_process_write: problems writing data ...\n");
  }

  /* update internal data */
  sion_filedesc->currentpos += bytestowrite;
  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)));

  return bwrote;
}
#undef DFUNCTION

bool _lsion_generic_collective_next_write_spec(void *iterator, int64_t *spec)
{
  _sion_range range;
  if (_lsion_partitioned_range_iterator_next((_sion_partitioned_range_iterator *)iterator, &range)) {
    _lsion_range_into_spec(range, spec);
    return true;
  } else {
    return false;
  }
}

#define DFUNCTION "_lsion_generic_collective_process_read"
int64_t _lsion_generic_collective_process_read(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_read: invalid sion_filedesc %d", sid);
  }
  DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "enter spec[0]=%" PRId64 " spec[1]=%" PRId64 " sid=%d\n", spec[0], spec[1], sid));

  /* move file pointer */
  int64_t destpos = spec[0], bytestoread = spec[1];
  if (sion_filedesc->currentpos != destpos) {
    if (sion_filedesc->fileptr != NULL) {
      _lsion_file_set_position(sion_filedesc->fileptr, destpos);
    }
    sion_filedesc->currentpos = destpos;
  }

  /* get and read data */
  DPRINTFP((4, DFUNCTION, _SION_DEFAULT_RANK, "collector start to read data of size %lld at position %lld\n",
    (long long)bytestoread, (long long)destpos));
  int64_t bread = _lsion_file_read(data, bytestoread, sion_filedesc->fileptr);
  if (bread != bytestoread) {
    return _lsion_errorprint(bread, _SION_ERROR_RETURN, "_lsion_generic_collective_process_read: problems reading data ...\n");
  }

  {
    ONLY_DEBUG(char *p = data;)
    DPRINTFP(
      (128, DFUNCTION, _SION_DEFAULT_RANK, "data[0]=%c data[%d]=%c\n", (char)p[0], (int)bytestoread, (char)p[bytestoread - 1]));
  }

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

  return bread;
}
#undef DFUNCTION

bool _lsion_generic_collective_next_read_spec(void *iterator, int64_t *spec)
{
  _sion_range range;
  if (_lsion_partitioned_limited_range_iterator_next((_sion_partitioned_limited_range_iterator *)iterator, &range)) {
    _lsion_range_into_spec(range, spec);
    return true;
  } else {
    return false;
  }
}

#define DFUNCTION "_lsion_coll_check_env"
int _lsion_coll_check_env(_sion_filedesc *sion_filedesc)
{
  const char *cs;
  const char *cn;
  const char *cd;
  int rc = SION_SUCCESS, numcoll;

  cd = _lsion_getenv("SION_COLLDEBUG");
  if (cd) {
    sion_filedesc->colldebug = atoi(cd);
  }

  cs = _lsion_getenv("SION_COLLSIZE");
  cn = _lsion_getenv("SION_COLLNUM");
  if (cs) {
    sion_filedesc->collsize = atoi(cs);
    if (sion_filedesc->collsize > sion_filedesc->ntasks) {
      sion_filedesc->collsize = sion_filedesc->ntasks;
    }
    if (sion_filedesc->colldebug >= 1) {
      fprintf(stderr, "collective statistics:            SION_COLLSIZE=%11d\n", sion_filedesc->collsize);
    }
  } else if (cn) {
    numcoll = atoi(cn);
    if (numcoll > 0) {
      if (numcoll > sion_filedesc->ntasks) {
        numcoll = sion_filedesc->ntasks;
      }
      sion_filedesc->collsize = sion_filedesc->ntasks / numcoll;
      if (sion_filedesc->ntasks % numcoll > 0) {
        sion_filedesc->collsize++;
      }

      if (sion_filedesc->colldebug >= 1) {
        fprintf(stderr, "collective statistics:             SION_COLLNUM=%11d\n", numcoll);
        fprintf(stderr, "collective statistics:                 collsize=%11d\n", sion_filedesc->collsize);
      }
    }
  }

  /* enable collective operation? */
  if ((cs) || (cn)) {
    if (sion_filedesc->collsize > 0) {
      sion_filedesc->usecoll = 1;
    }
    if (sion_filedesc->collsize < 0) {
      sion_filedesc->usecoll = 1;
    }
    if (sion_filedesc->collsize == 0) {
      sion_filedesc->usecoll = 0;
    }
  }

  DPRINTFP((2, DFUNCTION, _SION_DEFAULT_RANK, "usecoll=%d collsize=%d (%d tasks, %d files) colldebug=%d\n",
    sion_filedesc->usecoll, sion_filedesc->collsize, sion_filedesc->ntasks, sion_filedesc->nfiles, sion_filedesc->colldebug));

  return rc;
}
#undef DFUNCTION
