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

#include "sion.h"
#include "sion_debug.h"
#include "sion_error_handler.h"
#include "sion_fd.h"
#include "sion_file.h"
#include "sion_filedesc.h"
#include "sion_internal.h"

#define FILENAME_LENGTH 1024

static void usage(char *name);

int main(int argc, char **argv)
{
  _sion_fileptr *outfp;
  char *infilename;
  char outfilename[FILENAME_LENGTH];
  char *prefix;
  char fnmask[FILENAME_LENGTH];
  char fnkmask[FILENAME_LENGTH];

  int i, rank, blknum;
  char *localbuffer;
  int64_t chunksize = 0;
  int64_t left, bread, bsumread, bwrote;

  /* options */
  int digits = 5;
  int verbose = 0;
  int useglobalranks = 0;

  /* for file infomation */
  int sid, ntasks, nfiles, maxblocks;
  int32_t fsblksize;
  int64_t globalskip;
  int64_t start_of_varheader;
  int64_t *sion_localsizes;
  int64_t *sion_globalranks;
  int64_t *sion_blockcount;
  int64_t *sion_blocksizes;
  int64_t *sionsplit_sum_bytes_per_task;
  int64_t sionsplit_sum_bytes;
  int64_t sionsplit_filesize;

  _sion_filedesc *sion_filedesc;
  uint64_t key;

  /* parse command line */
  i = 1;
  if (argc < 3) {
    usage(argv[0]);
  }

  while (i < argc) {
    if (argv[i][0] == '-') {
      switch (argv[i][1]) {
      case 'd':
        digits = atoi(argv[++i]);
        break;
      case 'g':
        useglobalranks = 1;
        break;
      case 'v':
        verbose++;
        break;
      case 'V':
        fprintf(stderr, "SIONlib utility %s (Version %d.%dp%d, fileformat version %d)\n", argv[0], SION_MAIN_VERSION,
          SION_SUB_VERSION, SION_VERSION_PATCHLEVEL, SION_FILEFORMAT_VERSION);
        exit(1);
      case 'h':
        usage(argv[0]);
        break;
      default:
        usage(argv[0]);
      }
    }
    i++;
  }

  infilename = argv[argc - 2];
  prefix = argv[argc - 1];

  printf("sionsplit: filename:                  %-30s\n", infilename);
  printf("sionsplit: prefix:                    %-30s\n", prefix);

  sid = lsion_open(infilename, "rb,posix", &ntasks, &nfiles, NULL, &fsblksize, NULL, NULL);

  /* get the sion file structure */
  if ((_lsion_vcdtype(sid) != SION_FILEDESCRIPTOR) || !(sion_filedesc = _lsion_vcdtovcon(sid))) {
    return _lsion_errorprint(SION_NOT_SUCCESS, _SION_ERROR_RETURN, "invalid sion_filedesc, aborting %d ...\n", sid);
  }

  printf("sionsplit: sid:                       %d\n", sid);
  printf("sionsplit: filename:                  %-30s\n", infilename);
  printf("sionsplit: number of tasks:           %d\n", ntasks);
  printf("sionsplit: number of files:           %d\n", nfiles);
  printf("sionsplit: number of digits:          %d\n", digits);
  printf("sionsplit: use global ranks:          %d\n", useglobalranks);
  printf("sionsplit: keyval:                    %d (%s)\n", (int)sion_filedesc->keyvalmode,
    lsion_keyval_type_to_str(sion_filedesc->keyvalmode));
  printf("sionsplit: current endianness:        %s\n", (lsion_get_endianness()) ? "big" : "little");
  printf("sionsplit: file endianness:           %s\n", (lsion_get_file_endianness(sid)) ? "big" : "little");
  printf("sionsplit: fsblksize:                 %lu bytes (%6.2f MB)\n", (unsigned long)fsblksize, fsblksize / 1024.0 / 1024.0);

  lsion_get_locations(sid, &ntasks, &maxblocks, &globalskip, &start_of_varheader, &sion_localsizes, &sion_globalranks,
    &sion_blockcount, &sion_blocksizes);

  printf("sionsplit: max number of chunks:      %d\n", maxblocks);

  /* analysis */
  sionsplit_sum_bytes_per_task = malloc(ntasks * sizeof(int64_t));
  for (rank = 0; rank < ntasks; rank++) {
    sionsplit_sum_bytes_per_task[rank] = 0;
  }
  for (rank = 0; rank < ntasks; rank++) {
    for (blknum = 0; blknum < maxblocks; blknum++) {
      sionsplit_sum_bytes_per_task[rank] += sion_blocksizes[ntasks * blknum + rank];
    }
  }

  sionsplit_sum_bytes = 0;
  for (rank = 0; rank < ntasks; rank++) {
    sionsplit_sum_bytes += sionsplit_sum_bytes_per_task[rank];
  }

  printf("sionsplit: datasize in file (aggr.):  %" PRId64 " bytes  (%6.2f MB)\n", sionsplit_sum_bytes,
    sionsplit_sum_bytes / 1024.0 / 1024.0);

  sionsplit_filesize = start_of_varheader + (maxblocks + 1) * rank * sizeof(int64_t);

  printf("sionsplit: start_of_varheader:        %" PRId64 " bytes  (%6.2f MB)\n", start_of_varheader,
    start_of_varheader / 1024.0 / 1024.0);
  printf("sionsplit: size of file:              %" PRId64 " bytes  (%6.2f MB)\n", sionsplit_filesize,
    sionsplit_filesize / 1024.0 / 1024.0);

  if (sionsplit_filesize > 0) {
    printf("sionsplit: file usage:                %8.6f%%\n", (double)sionsplit_sum_bytes / (double)sionsplit_filesize * 100.0);
  }
  chunksize = 0;

  if (FILENAME_LENGTH <= snprintf(fnmask, FILENAME_LENGTH, "%s%%0%dd", prefix, digits)) {
    return _lsion_errorprint(SION_NOT_SUCCESS, _SION_ERROR_RETURN, "prefix is too long, aborting: %s\n", prefix);
  }
  printf("sionsplit: filename mask:             %-s\n", fnmask);

  if (FILENAME_LENGTH <= snprintf(fnkmask, FILENAME_LENGTH, "%s%%0%dd_%%012ld", prefix, digits)) {
    return _lsion_errorprint(SION_NOT_SUCCESS, _SION_ERROR_RETURN, "prefix is too long, aborting: %s\n", prefix);
  }
  printf("sionsplit: filename key mask:         %-s\n", fnkmask);

  if (sion_filedesc->keyvalmode == SION_KEYVAL_NONE) {
    /* standard copy with read/write */

    for (rank = 0; rank < ntasks; rank++) {
      if (chunksize < sion_localsizes[rank]) {
        chunksize = sion_localsizes[rank];
      }
      if (chunksize < sion_blocksizes[rank]) {
        chunksize = sion_blocksizes[rank];
      }
    }
    printf("sionsplit: max chunksize:             %" PRId64 "\n", chunksize);
    localbuffer = malloc(chunksize);
    if (localbuffer == NULL) {
      fprintf(stderr, "cannot allocate localbuffer of size %" PRId64 " , aborting ...\n", chunksize);
      free(sionsplit_sum_bytes_per_task);
      return 1;
    }

    for (rank = 0; rank < ntasks; rank++) {
      if (FILENAME_LENGTH
        <= snprintf(outfilename, FILENAME_LENGTH, fnmask, (useglobalranks) ? (int)sion_globalranks[rank] : rank)) {
        fprintf(stderr, "output file name too long, aborting: ");
        fprintf(stderr, fnmask, (useglobalranks) ? (int)sion_globalranks[rank] : rank);
        fprintf(stderr, "\n");
        return 1;
      }

      printf("sionsplit: generating file:           %-s\n", outfilename);
      outfp = _lsion_file_open(outfilename, SION_FILE_FLAG_ANSI | SION_FILE_FLAG_WRITE | SION_FILE_FLAG_CREATE, 0);
      if (outfp == NULL) {
        fprintf(stderr, "cannot open outfile %s , aborting ...\n", outfilename);
        free(sionsplit_sum_bytes_per_task);
        free(localbuffer);
        return 1;
      }

      for (blknum = 0; blknum < sion_blockcount[rank]; blknum++) {
        /* seek position of block */
        lsion_switch_logical_file(sid, rank);
        lsion_seek_chunk(sid, blknum, SION_SEEK_BEGIN);
        DPRINTFP((1, "sionsplit", 0, "after lsion_seek sid=%d rank=%d blknum=%d fileposition=%" PRId64 "\n", sid, rank, blknum,
          lsion_get_position(sid)));

        /* read data from block */
        left = sion_blocksizes[ntasks * blknum + rank];
        bsumread = 0;
        while (left > 0) {
          DPRINTFP((8, "sionsplit", 0, "will read %" PRId64 " bytes localbuffer+%" PRId64 "\n", left, bsumread));
          bread = lsion_fread(localbuffer + bsumread, 1, left, sid);
          left -= bread;
          bsumread += bread;
          printf("sionsplit:                            %" PRId64 " read left=%" PRId64 " \n", bread, left);
        }

        /* write data to outfile */
        left = sion_blocksizes[ntasks * blknum + rank];
        bwrote = _lsion_file_write(localbuffer, left, outfp);
        printf("sionsplit:                            %" PRId64 " wrote of left=%" PRId64 " \n", bwrote, left);
      }

      _lsion_file_close(outfp);
    }

    free(localbuffer);

  } else {
    /* key-value copy with read/write */

    printf("sionsplit: buffer size:             %d\n", fsblksize);
    localbuffer = malloc(fsblksize);
    if (localbuffer == NULL) {
      fprintf(stderr, "cannot allocate localbuffer of size %" PRId64 " , aborting ...\n", chunksize);
      free(sionsplit_sum_bytes_per_task);
      return 1;
    }

    for (rank = 0; rank < ntasks; rank++) {
      if (verbose) {
        printf("siondefrag: ->rank:                    %d\n", rank);
      }
      if (!verbose)
        if (rank % 16 == 0) {
          printf("[%d]", rank);
          fflush(stdout);
        }

      /* search rank in input file */
      lsion_switch_logical_file(sid, rank);

      lsion_key_full_scan(sid);

      /* reset iterator over keys */
      lsion_key_list_iterator_reset(sid);

      /* loop over key-value blocks */
      while (lsion_key_list_iterator_next(sid, &key) == SION_SUCCESS) {
        if (FILENAME_LENGTH
          <= snprintf(outfilename, FILENAME_LENGTH, fnkmask, (useglobalranks) ? (int)sion_globalranks[rank] : rank, (long)key)) {
          fprintf(stderr, "output file name too long, aborting: ");
          fprintf(stderr, fnkmask, (useglobalranks) ? (int)sion_globalranks[rank] : rank, (long)key);
          fprintf(stderr, "\n");
          return 1;
        }

        printf("sionsplit: generating file:           '%-s'\n", outfilename);
        outfp = _lsion_file_open(outfilename, SION_FILE_FLAG_ANSI | SION_FILE_FLAG_WRITE | SION_FILE_FLAG_CREATE, 0);
        if (outfp == NULL) {
          fprintf(stderr, "cannot open outfile %s , aborting ...\n", outfilename);
          free(localbuffer);
          free(sionsplit_sum_bytes_per_task);
          return 1;
        }

        while ((bread = lsion_fread_key(localbuffer, key, 1, fsblksize, sid)) > 0) {
          if (bread > 0) {
            bwrote = _lsion_file_write(localbuffer, bread, outfp);
          }
          if (verbose)
            printf("sionsplit:                               extracting now data of key[%12ld] (%" PRId64 " bytes) (%" PRId64
                   " bytes written)\n",
              (long)key, bread, bwrote);
        }

        _lsion_file_close(outfp);
      }
    }
    free(localbuffer);
  }

  free(sionsplit_sum_bytes_per_task);
  lsion_close(sid);

  return 0;
}

void usage(char *name)
{
  fprintf(stderr, "Usage: %s options <sionfn> <prefix>\n\n", name);

  fprintf(stderr, "Split SIONlib file <sionfn> into separate files. For each task a file is\n");
  fprintf(stderr, "created with the name <prefix><digits>\n\n");

  fprintf(stderr, "Example: %s data.sion data/file_\n", name);
  fprintf(stderr, "  creates files data/file_00000, data/file_00001, ...\n\n");

  fprintf(stderr, "Options:\n");
  fprintf(stderr, "  [-v]         verbose mode\n");
  fprintf(stderr, "  [-g]         use global rank for numbering files\n");
  fprintf(stderr, "  [-d <num>]   number of digits for filename generation (default 5)\n");
  fprintf(stderr, "  [-V]         show version of SIONlib\n");
  fprintf(stderr, "  [-h]         show this help\n");
  exit(1);
}
