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

#ifdef SION_MPI

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

#include "sion_const.h"
#include "sion_debug.h"
#include "sion_enums.h"
#include "sion_error_handler.h"
#include "sion_flags.h"
#include "sion_generic.h"
#include "sion_generic_internal.h"
#include "sion_internal.h"
#include "sion_mpi.h"
#include "sion_mpi_cb_gen.h"
#include "sion_mpi_internal_gen.h"

int _lsion_mpi_api_aid = -1;

int lsion_paropen_mpi(const char *fname, const char *file_mode, int *numFiles, MPI_Comm gComm, const MPI_Comm *lComm,
  int64_t *chunksize, int32_t *fsblksize, int *globalrank, FILE **fileptr, char **newfname)
{
  /* parse file mode */
  int sid;
  _sion_flags_store *flags_store = _lsion_parse_flags(file_mode);
  if (!flags_store) {
    return _lsion_errorprint_mpi(
      SION_ID_NOT_VALID, _SION_ERROR_RETURN, "lsion_paropen_mpi: could not parse file mode in %s, aborting ...\n", file_mode);
  }
  sid = _lsion_paropen_mpi(fname, flags_store, numFiles, gComm, lComm, chunksize, fsblksize, globalrank, fileptr, newfname);
  _lsion_flags_destroy_store(&flags_store);
  return sid;
}

int _lsion_paropen_mpi(const char *fname, const _sion_flags_store *flags_store, int *numFiles, MPI_Comm gComm,
  const MPI_Comm *lComm, int64_t *chunksize, int32_t *fsblksize, int *globalrank, FILE **fileptr, char **newfname)
{
  int rc = SION_NOT_SUCCESS, sid = SION_ID_UNDEF;
  int filenumber, gtasks, gRank, lRank, lSize;

  _mpi_api_commdata *gen_gcomm;

  MPI_Comm_size(gComm, &gtasks);
  MPI_Comm_rank(gComm, &gRank);

  DPRINTFP((1, "lsion_paropen_mpi", gRank, "enter parallel open of file %s\n", fname));

  /* check parameters */
  if (lComm == NULL) {
    return _lsion_errorprint_mpi(SION_ID_NOT_VALID, _SION_ERROR_RETURN, "lsion_paropen_mpi: No lComm variable given");
  }
  if (numFiles == NULL) {
    return _lsion_errorprint_mpi(SION_ID_NOT_VALID, _SION_ERROR_RETURN, "lsion_paropen_mpi: No numFiles variable given");
  }

  /* register callbacks for generic interface */
  if (_lsion_mpi_api_aid < 0) {
    _lsion_mpi_api_aid = _lsion_register_callbacks_mpi();
  }

  if (flags_store->mask & _SION_FMODE_WRITE) {
    /* file mode WRITE */

    if (*numFiles <= 0) {
      /* lComm contains local communicator */

      rc = _lsion_get_info_from_splitted_comm_mpi(gComm, *lComm, numFiles, &filenumber, &lRank, &lSize);
      if (rc != SION_SUCCESS) {
        return _lsion_errorprint_mpi(
          SION_ID_NOT_VALID, _SION_ERROR_RETURN, "lsion_paropen_mpi: error in _lsion_get_info_from_splitted_comm_mpi");
      }
      DPRINTFP((1, "lsion_paropen_mpi", gRank, "%d local communicators found\n", *numFiles));

    } else {
      /* number of files is given */

      rc = _lsion_gen_info_from_gcomm_mpi(*numFiles, gComm, &filenumber, &lRank, &lSize);
      if (rc != SION_SUCCESS) {
        return _lsion_errorprint_mpi(
          SION_ID_NOT_VALID, _SION_ERROR_RETURN, "lsion_paropen_mpi: error in _lsion_gen_info_from_gcomm_mpi");
      }
      DPRINTFP((1, "lsion_paropen_mpi", gRank, "Global communicator divided in %d local communicators\n", *numFiles));
    }

    /* overwrite globalrank set by user, necessary for multi-file support */
    *globalrank = gRank;

  } else if (flags_store->mask & _SION_FMODE_READ) {
    /* file mode READ */
    /* nothing to do info will be returned by generic paropen */

    /* set to gRank, current rank in global communicator, this is
       different to older versions of SIONlib, where globalrank comes
       from file in read case */
    *globalrank = gRank;

    if (!(flags_store->mask & _SION_FMODE_BUDDY)) {
      lRank = lSize = -1; /* will be set by lsion_generic_paropen */
    } else {
      /* lvomm must be given for buddy checkpointing */
      rc = _lsion_get_info_from_splitted_comm_mpi(gComm, *lComm, numFiles, &filenumber, &lRank, &lSize);
      if (rc != SION_SUCCESS) {
        return _lsion_errorprint_mpi(
          SION_ID_NOT_VALID, _SION_ERROR_RETURN, "lsion_paropen_mpi: error in _lsion_get_info_from_splitted_comm_mpi");
      }
      DPRINTFP((1, "lsion_paropen_mpi", gRank, "%d local communicators found\n", *numFiles));
    }

  } else {
    return _lsion_errorprint_mpi(SION_ID_NOT_VALID, _SION_ERROR_RETURN, "lsion_paropen_mpi: unknown file mode");
  }

  /* create generic communicator container */
  gen_gcomm = malloc(sizeof(_mpi_api_commdata));
  if (gen_gcomm == NULL) {
    return _lsion_errorprint(SION_ID_NOT_VALID, _SION_ERROR_RETURN,
      "cannot allocate mpi internal data structure of size %lu (_mpi_api_commdata), aborting ...\n",
      (unsigned long)sizeof(_mpi_api_commdata));
  }
  gen_gcomm->comm = gComm;
  gen_gcomm->commset = 1;
  gen_gcomm->local = 0;
  gen_gcomm->rank = gRank;
  gen_gcomm->size = gtasks;
  gen_gcomm->lcommgroup = NULL;

  /* FIXME: add debug output for flags_store
  DPRINTFP((1, "lsion_paropen_mpi", gRank, "enter parallel open of %d files (current name %s) in %s mode\n", *numFiles, fname, file_mode)); */
  DPRINTFP((1, "lsion_paropen_mpi", gRank, "enter parallel open of %d files (current name %s)\n", *numFiles, fname));
  sid = _lsion_generic_paropen(_lsion_mpi_api_aid, fname, flags_store, chunksize, fsblksize, gen_gcomm, gRank, gtasks, &filenumber,
    numFiles, &lRank, &lSize, fileptr, newfname);
  /* FIXME: add debug output for flags_store
  DPRINTFP((1, "lsion_paropen_mpi", gRank, "leave parallel open of %d files in %s mode #tasks=%d sid=%d\n", *numFiles, file_mode, lSize, sid)); */
  DPRINTFP((1, "lsion_paropen_mpi", gRank, "leave parallel open of %d files #tasks=%d sid=%d\n", *numFiles, lSize, sid));

  /* test return code from internal open */
  if (sid == SION_ID_NOT_VALID) {
    return _lsion_errorprint_mpi(
      SION_ID_NOT_VALID, _SION_ERROR_RETURN, "lsion_paropen_mpi: invalid return code from internal open %d", rc);
  }

  DPRINTFP((1, "lsion_paropen_mpi", gRank, "leave parallel open of file %s sid=%d\n", fname, sid));

  return sid;
}

int lsion_parclose_mpi(int sid)
{
  int rc = 0;

  DPRINTFP((1, "lsion_parclose_mpi", _SION_DEFAULT_RANK, "enter parallel close of sid %d\n", sid));

  rc = lsion_generic_parclose(sid);

  DPRINTFP((1, "lsion_parclose_mpi", _SION_DEFAULT_RANK, "leave parallel close of sid %d rc=%d\n", sid, rc));

  return rc;
}

int lsion_parreinit_mpi(int sid, int64_t chunksize)
{
  int rc = 0;

  DPRINTFP((1, "lsion_parreinit_mpi", _SION_DEFAULT_RANK, "enter parallel reinit of sid %d\n", sid));

  rc = lsion_generic_parreinit(sid, chunksize);

  DPRINTFP((1, "lsion_parreinit_mpi", _SION_DEFAULT_RANK, "leave parallel reinit of sid %d rc=%d\n", sid, rc));

  return rc;
}

int lsion_paropen_mapped_mpi(char *fname, const char *file_mode, int *numFiles, MPI_Comm gComm, int *nlocaltasks,
  int **globalranks, int64_t **chunksizes, int **mapping_filenrs, int **mapping_lranks, int32_t *fsblksize, FILE **fileptr)
{
  int sid;
  _sion_flags_store *flags_store = _lsion_parse_flags(file_mode);
  if (!flags_store) {
    return _lsion_errorprint_mpi(SION_ID_NOT_VALID, _SION_ERROR_RETURN,
      "lsion_paropen_mapped_mpi: could not parse file mode in %s, aborting ...\n", file_mode);
  }
  sid = _lsion_paropen_mapped_mpi(fname, flags_store, numFiles, gComm, nlocaltasks, globalranks, chunksizes, mapping_filenrs,
    mapping_lranks, fsblksize, fileptr);
  _lsion_flags_destroy_store(&flags_store);
  return sid;
}

int _lsion_paropen_mapped_mpi(char *fname, const _sion_flags_store *flags_store, int *numFiles, MPI_Comm gComm, int *nlocaltasks,
  int **globalranks, int64_t **chunksizes, int **mapping_filenrs, int **mapping_lranks, int32_t *fsblksize, FILE **fileptr)
{
  int sid = SION_ID_UNDEF;
  int gtasks, gRank;
  _mpi_api_commdata *gen_gcomm;

  MPI_Comm_size(gComm, &gtasks);
  MPI_Comm_rank(gComm, &gRank);

  DPRINTFP((1, "lsion_paropen_mapped_mpi", gRank, "enter parallel open of file %s\n", fname));

  /* check parameters */
  if (numFiles == NULL) {
    return _lsion_errorprint_mpi(SION_ID_NOT_VALID, _SION_ERROR_RETURN, "lsion_paropen_mapped_mpi: No numFiles variable given");
  }

  /* register callbacks for generic interface */
  if (_lsion_mpi_api_aid < 0) {
    _lsion_mpi_api_aid = _lsion_register_callbacks_mpi();
  }

  if (flags_store->mask & _SION_FMODE_WRITE) {
    /* file mode WRITE */

    if (*numFiles <= 0) {
      return _lsion_errorprint_mpi(SION_ID_NOT_VALID, _SION_ERROR_RETURN,
        "lsion_paropen_mapped_mpi: numFiles variable <= 0 not allowed for mapped files in write mode");
    }
  } else if (flags_store->mask & _SION_FMODE_READ) {
    /* file mode READ */
    /* nothing to do here so far, filenumbers and mapping will be determined by in generic routine */

  } else {
    return _lsion_errorprint_mpi(SION_ID_NOT_VALID, _SION_ERROR_RETURN, "lsion_paropen_mapped_mpi: unknown file mode");
  }

  /* create generic communicator container */
  gen_gcomm = malloc(sizeof(_mpi_api_commdata));
  if (gen_gcomm == NULL) {
    return _lsion_errorprint(SION_ID_NOT_VALID, _SION_ERROR_RETURN,
      "cannot allocate mpi internal data structure of size %lu (_mpi_api_commdata), aborting ...\n",
      (unsigned long)sizeof(_mpi_api_commdata));
  }
  gen_gcomm->comm = gComm;
  gen_gcomm->commset = 1;
  gen_gcomm->local = 0;
  gen_gcomm->rank = gRank;
  gen_gcomm->size = gtasks;
  gen_gcomm->lcommgroup = NULL;

  // FIXME: re-enable output of file mode
  // DPRINTFP((1, "lsion_paropen_mapped_mpi", gRank, "enter parallel open of %d files (current name %s) in %s mode (sid=%d)\n", *numFiles, fname, file_mode, sid));
  DPRINTFP(
    (1, "lsion_paropen_mapped_mpi", gRank, "enter parallel open of %d files (current name %s) (sid=%d)\n", *numFiles, fname, sid));
  sid = _lsion_generic_paropen_mapped(_lsion_mpi_api_aid, fname, flags_store, numFiles, gen_gcomm, gRank, gtasks, nlocaltasks,
    globalranks, chunksizes, mapping_filenrs, mapping_lranks, fsblksize, fileptr);
  // FIXME: re-enable output of file mode
  // DPRINTFP((1, "lsion_paropen_mapped_mpi", gRank, "leave parallel open of %d files in %s mode #tasks=%d sid=%d\n", *numFiles, file_mode, *nlocaltasks, sid));
  DPRINTFP(
    (1, "lsion_paropen_mapped_mpi", gRank, "leave parallel open of %d files #tasks=%d sid=%d\n", *numFiles, *nlocaltasks, sid));

  /* test return code from internal open */
  if (sid == SION_ID_NOT_VALID) {
    return _lsion_errorprint_mpi(
      SION_ID_NOT_VALID, _SION_ERROR_RETURN, "lsion_paropen_mapped_mpi: invalid return code from internal open %d", sid);
  }

  DPRINTFP((1, "lsion_paropen_mapped_mpi", gRank, "leave parallel open of file %s sid=%d\n", fname, sid));

  return sid;
}

int lsion_parclose_mapped_mpi(int sid)
{
  int rc = 0;

  DPRINTFP((1, "lsion_parclose_mapped_mpi", _SION_DEFAULT_RANK, "enter parallel close of sid %d\n", sid));

  rc = lsion_generic_parclose_mapped(sid);

  DPRINTFP((1, "lsion_parclose_mapped_mpi", _SION_DEFAULT_RANK, "leave parallel close of sid %d rc=%d\n", sid, rc));

  return rc;
}

sion_mpi_options lsion_mpi_options_new()
{
  sion_mpi_options options;

  options.chunksize = 1024 * 1024; // FIXME: default to fsblksize (which is only known on open)
  options.fsblksize = -1;

  options.multifile_mode = _SION_MULTIFILE_SINGLE;

  options.keyval_mode = SION_KEYVAL_MODE_NONE;

  options.buddylevel = 0;

  options.collsize = 0;
  options.collective_merge = false;

  return options;
}

void lsion_mpi_options_set_chunksize(sion_mpi_options *options, int64_t chunksize)
{
  // TODO: assert write mode
  options->chunksize = chunksize;
}

void lsion_mpi_options_set_fsblksize(sion_mpi_options *options, int32_t fsblksize)
{
  // TODO: assert write mode
  // TODO: adapt chunksize?
  options->fsblksize = fsblksize;
}

void lsion_mpi_options_set_multifile_number(sion_mpi_options *options, int multifile_number)
{
  // TODO: assert write mode
  options->multifile_mode = _SION_MULTIFILE_NUMBER;
  options->multifile.number = multifile_number;
}

void lsion_mpi_options_set_multifile_communicator(sion_mpi_options *options, MPI_Comm multifile_communicator)
{
  options->multifile_mode = _SION_MULTIFILE_COMMUNICATOR;
  options->multifile.communicator = multifile_communicator;
}

void lsion_mpi_options_set_keyval_mode(sion_mpi_options *options, sion_keyval_mode keyval_mode)
{
  options->keyval_mode = keyval_mode;
}

void lsion_mpi_options_set_buddy(sion_mpi_options *options)
{
  lsion_mpi_options_set_buddylevel(options, 1);
}

void lsion_mpi_options_set_buddylevel(sion_mpi_options *options, int32_t buddylevel)
{
  options->buddylevel = buddylevel;
}

void lsion_mpi_options_set_collective(sion_mpi_options *options)
{
  options->collsize = -1;
}

void lsion_mpi_options_set_collective_size(sion_mpi_options *options, int32_t size)
{
  options->collsize = size;
}

void lsion_mpi_options_set_collective_merge(sion_mpi_options *options)
{
  if (options->collsize == 0) {
    options->collsize = -1;
  }
  options->collective_merge = true;
}

_sion_flags_store *_lsion_mpi_options_into_flags_store(sion_open_mode mode, const sion_mpi_options *options)
{
  _sion_flags_store *flags_store = _lsion_flags_create_store();
  if (!flags_store) {
    /* FIXME: add error message */
    fprintf(stderr, "could not create flags store.\n");
    return NULL;
  }

  switch (mode) {
  case SION_OPEN_READ:
    _lsion_flags_add(flags_store, "br", "");

    switch (options->keyval_mode) {
    case SION_KEYVAL_MODE_NONE:
      /* do nothing */
      break;
    case SION_KEYVAL_MODE_DEFAULT:
    /* FALLTHROUGH */
    case SION_KEYVAL_MODE_INLINE:
      _lsion_flags_add(flags_store, "keyval", "inline");
      break;
    case SION_KEYVAL_MODE_META:
      _lsion_flags_add(flags_store, "keyval", "meta");
      break;
    case SION_KEYVAL_MODE_HASH:
      _lsion_flags_add(flags_store, "keyval", "hash");
      break;
    case SION_KEYVAL_MODE_UNKNOWN:
      _lsion_flags_add(flags_store, "keyval", "unknown");
      break;
    default:
      /* FIXME: add error message */
      fprintf(stderr, "wrong keyval mode in read mode.\n");
      _lsion_flags_destroy_store(&flags_store);
      return NULL;
    }

    break;
  case SION_OPEN_WRITE:
    _lsion_flags_add(flags_store, "bw", "");

    switch (options->keyval_mode) {
    case SION_KEYVAL_MODE_NONE:
      /* do nothing */
      break;
    case SION_KEYVAL_MODE_DEFAULT:
    /* FALLTHROUGH */
    case SION_KEYVAL_MODE_INLINE:
      _lsion_flags_add(flags_store, "keyval", "inline");
      break;
    case SION_KEYVAL_MODE_META:
      _lsion_flags_add(flags_store, "keyval", "meta");
      break;
    case SION_KEYVAL_MODE_HASH:
      _lsion_flags_add(flags_store, "keyval", "hash");
      break;
    default: /* 'unknown' keyval mode is not allowed in write mode */
      /* FIXME: add error message */
      fprintf(stderr, "wrong keyval mode in read mode.\n");
      _lsion_flags_destroy_store(&flags_store);
      return NULL;
    }

    break;
  }

  if (options->buddylevel > 0) {
    char buddylevel_str[12];
    int written = snprintf(buddylevel_str, 12, "%" PRId32, (int32_t)options->buddylevel);
    // FIXME: assert written < 12
    (void)written;
    _lsion_flags_add(flags_store, "buddy", buddylevel_str);
    fprintf(stderr, "added buddy level %s.\n", buddylevel_str);
  }

  if (options->collsize) {
    _lsion_flags_add(flags_store, "collective", "");
    char collsize_str[12];
    int written = snprintf(collsize_str, 12, "%" PRId32, (int32_t)options->collsize);
    // FIXME: assert written < 12
    (void)written;
    _lsion_flags_add(flags_store, "collsize", collsize_str);
    if (options->collective_merge) {
      _lsion_flags_add(flags_store, "cmerge", "");
    }
  }

  _lsion_flags_update_mask(flags_store);

  return flags_store;
}

int lsion_paropen_mpi_with_options(
  const char *filename, sion_open_mode mode, MPI_Comm communicator, const sion_mpi_options *options_)
{
  sion_mpi_options options; // FIXME: need local copy to satisfy signature of lsion_paropen_mpi()
  if (options_ == NULL) {
    options = lsion_mpi_options_new();
  } else {
    options = *options_;
  }

  int num_files;
  MPI_Comm lcomm;
  int rank_;

  switch (options.multifile_mode) {
  case _SION_MULTIFILE_SINGLE:
    num_files = 1;
    lcomm = communicator;
    break;
  case _SION_MULTIFILE_NUMBER:
    num_files = options.multifile.number;
    lcomm = communicator;
    break;
  case _SION_MULTIFILE_COMMUNICATOR:
    num_files = -1;
    lcomm = options.multifile.communicator;
    break;
  }

  _sion_flags_store *flags_store = _lsion_mpi_options_into_flags_store(mode, &options);
  if (!flags_store) {
    return -1;
  }

  int sid = _lsion_paropen_mpi(
    filename, flags_store, &num_files, communicator, &lcomm, &options.chunksize, &options.fsblksize, &rank_, NULL, NULL);

  _lsion_flags_destroy_store(&flags_store);
  return sid;
}

/* end of ifdef MPI */
#endif
