/****************************************************************************
**  SIONLIB     http://www.fz-juelich.de/jsc/sionlib                       **
*****************************************************************************
**  Copyright (c) 2008-2019                                                **
**  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_error_handler.h"
#include "sion_fd.h"
#include "sion_filedesc.h"
#include "sion_internal.h"

#define FILENAME_LENGTH 1024
#define FILE_MODE_LENGTH 1024
#define MAXNULLCNT 1000

#define MB *1024 * 1024

static void usage(char *name);

int main(int argc, char **argv)
{
  char infilename[FILENAME_LENGTH];
  char outfilename[FILENAME_LENGTH];

  int i, rank, blknum;
  int nullcount;
  int64_t max_chunksize = 0;
  char *localbuffer = NULL;
  long localbuffer_size = -1;
  int64_t left, bread, bsumread, bwrote, bsumwrote;

  /* options */
  int verbose = 0;
  /* int64_t opt_localsize = -1; */
  int32_t opt_fsblocksize = -1;

  /* for file infomation */
  int ntasks, nfiles;
  int32_t fsblksize;
  int sid, outsid, size, blocks;
  int64_t globalskip;
  int64_t start_of_varheader;
  int64_t *sion_chunksizes;
  int64_t *sion_globalranks;
  int64_t *sion_blockcount;
  int64_t *sion_blocksizes;
  int64_t *siondefrag_sum_bytes_per_task;
  int64_t siondefrag_sum_bytes;
  int64_t siondefrag_filesize;
  int64_t siondefrag_sum_chunksizes;
  _sion_filedesc *sion_filedesc;
  uint64_t key;
  size_t len;

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

  while (i < argc) {
    if (argv[i][0] == '-') {
      switch (argv[i][1]) {
      /*
      case 's':
      opt_localsize = atoi(argv[++i]);
      break;
      case 'S':
      opt_localsize = atoi(argv[++i]) MB;
      break;
      */
      case 'q':
        opt_fsblocksize = atoi(argv[++i]);
        break;
      case 'Q':
        opt_fsblocksize = atoi(argv[++i]) MB;
        break;
      case 'v':
        verbose++;
        break;
      case 'V':
        fprintf(stderr, "SIONlib utility %s (Version %d.%dp%d, fileformat version %d)\n", argv[0], SION_VERSION_MAJOR,
          SION_VERSION_MINOR, SION_VERSION_PATCH, SION_FILEFORMAT_VERSION);
        exit(1);
      case 'h':
        usage(argv[0]);
        break;
      default:
        usage(argv[0]);
      }
    }
    i++;
  }

  strcpy(infilename, argv[argc - 2]);
  strcpy(outfilename, argv[argc - 1]);

  printf("siondefrag: infilename:                %-30s\n", infilename);
  printf("siondefrag: outfilename:               %-30s\n", outfilename);

  sid = sion_open(infilename, SION_OPEN_READ, 0, NULL);
  ntasks = sion_get_number_of_logical_files(sid);
  nfiles = sion_get_number_of_files(sid);
  fsblksize = sion_get_fsblksize(sid);

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

  printf("siondefrag: sid:                       %d\n", sid);
  printf("siondefrag: filename:                  %-30s\n", infilename);
  printf("siondefrag: nfiles:                    %d\n", nfiles);
  printf("siondefrag: fsblksize:                 %lu bytes (%6.2f MB)\n", (unsigned long)fsblksize, fsblksize / 1024.0 / 1024.0);
  printf("siondefrag: keyval:                    %d (%s)\n", (int)sion_filedesc->keyvalmode,
    sion_keyval_type_to_str(sion_filedesc->keyvalmode));
  printf("siondefrag: current endianness:        %s\n", (sion_get_endianness()) ? "big" : "little");
  printf("siondefrag: file endianness:           %s\n", (sion_get_file_endianness(sid)) ? "big" : "little");

  sion_get_locations(sid, &size, &blocks, &globalskip, &start_of_varheader, &sion_chunksizes, &sion_globalranks, &sion_blockcount,
    &sion_blocksizes);

  printf("siondefrag: max number of blocks:      %d\n", blocks);

  /* analysis */
  siondefrag_sum_bytes_per_task = malloc(ntasks * sizeof(int64_t));
  for (rank = 0; rank < size; rank++) {
    siondefrag_sum_bytes_per_task[rank] = 0;
  }
  for (rank = 0; rank < size; rank++) {
    for (blknum = 0; blknum < sion_blockcount[rank]; blknum++) {
      siondefrag_sum_bytes_per_task[rank] += sion_blocksizes[size * blknum + rank];
    }
  }
  siondefrag_sum_chunksizes = 0;
  siondefrag_sum_bytes = 0;
  for (rank = 0; rank < size; rank++) {
    siondefrag_sum_bytes += siondefrag_sum_bytes_per_task[rank];
    siondefrag_sum_chunksizes += sion_chunksizes[rank];
  }
  printf("siondefrag: datasize in file (aggr.):  %" PRId64 " bytes  (%6.2f MB)\n", siondefrag_sum_bytes,
    siondefrag_sum_bytes / 1024.0 / 1024.0);
  siondefrag_filesize = start_of_varheader + (blocks + 1) * rank * sizeof(int64_t);
  printf("siondefrag: start_of_varheader:        %" PRId64 " bytes  (%6.2f MB)\n", start_of_varheader,
    start_of_varheader / 1024.0 / 1024.0);
  printf("siondefrag: size of chunks:            %" PRId64 " bytes  (%6.2f MB)\n", siondefrag_sum_chunksizes,
    siondefrag_sum_chunksizes / 1024.0 / 1024.0);
  if (siondefrag_filesize > 0) {
    printf("siondefrag: file usage:                %8.6f%%\n",
      (double)siondefrag_sum_bytes / (double)siondefrag_sum_chunksizes * 100.0);
  }

  /* scan for block with max. number of bytes  */
  for (rank = 0; rank < size; rank++) {
    if (sion_chunksizes[rank] > max_chunksize) {
      max_chunksize = sion_chunksizes[rank];
    }
    for (blknum = 0; blknum < sion_blockcount[rank]; blknum++) {
      if (sion_blocksizes[size * blknum + rank] > max_chunksize) {
        max_chunksize = sion_blocksizes[size * blknum + rank];
      }
    }
  }

  if (opt_fsblocksize != -1) {
    fsblksize = opt_fsblocksize;
  }

  if ((sion_filedesc->keyvalmode != SION_KEYVAL_NONE) && (sion_filedesc->keyvalmode != SION_KEYVAL_INLINE)) {
    fprintf(stderr, "unknown keyvalue mode ..., aborting\n");
    free(siondefrag_sum_bytes_per_task);
    return 1;
  }

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

    printf("siondefrag: max chunksize:             %" PRId64 "\n", max_chunksize);
    localbuffer = malloc(max_chunksize);
    if (localbuffer == NULL) {
      free(siondefrag_sum_bytes_per_task);
      fprintf(stderr, "cannot allocate localbuffer of size %" PRId64 " , aborting ...\n", max_chunksize);
      return 1;
    }

    sion_options *options = sion_options_new();
    if (!options) {
      fprintf(stderr, "siondefrag: could not allocate options, aborting\n");
      return EXIT_FAILURE;
    }
    sion_options_set_lowlevel_api(options, SION_LOWLEVEL_API_POSIX);
    sion_options_set_endianness(options, sion_get_file_endianness(sid) ? SION_ENDIANNESS_BIG : SION_ENDIANNESS_LITTLE);
    sion_options_set_chunksizes(options, ntasks, siondefrag_sum_bytes_per_task);
    sion_options_set_fsblksize(options, fsblksize);
    outsid = sion_open(outfilename, SION_OPEN_WRITE, ntasks, options);
    sion_options_delete(options);

    printf("siondefrag: outsid:                    %d\n", outsid);
    printf("siondefrag: fsblksize outfile:         %lu\n", (unsigned long)fsblksize);

    for (rank = 0; rank < size; rank++) {
      /* set position in outfile to current rank */
      if (verbose) {
        printf("siondefrag: ->rank:                    %d\n", rank);
      }
      if (!verbose)
        if (rank % 16 == 0) {
          printf("[%d]", rank);
          fflush(stdout);
        }
      sion_switch_logical_file(outsid, rank);

      if (verbose) {
        printf("siondefrag:                            copy now %" PRId64 " blocks on rank=%d \n", sion_blockcount[rank], rank);
      }
      for (blknum = 0; blknum < sion_blockcount[rank]; blknum++) {
        /* seek position of block */
        sion_seek_chunk(sid, blknum, SION_SEEK_BEGINNING);

        /* read data from block */
        left = sion_blocksizes[size * blknum + rank];
        if (verbose) {
          printf("siondefrag:                               copy now block #%d with %" PRId64 " bytes \n", blknum, left);
        }
        bsumread = 0;
        nullcount = 0;
        while (left > 0) {
          bread = sion_fread(localbuffer + bsumread, 1, left, sid);
          if (bread == 0) {
            nullcount++;
          } else {
            nullcount = 0;
          }
          left -= bread;
          bsumread += bread;
          /*         printf("siondefrag:                            %" PRId64 " read left=%" PRId64 " \n", bread, left); */
          if (nullcount > MAXNULLCNT) {
            fprintf(stderr, "timeout on read data , aborting ...\n");
            exit(0);
          }
        }

        /* write data to outfile */
        left = sion_blocksizes[size * blknum + rank];
        bsumwrote = 0;
        while (left > 0) {
          bwrote = sion_fwrite(localbuffer + bsumwrote, 1, left, outsid);
          if (bwrote == 0) {
            nullcount++;
          } else {
            nullcount = 0;
          }
          left -= bwrote;
          bsumwrote += bwrote;
          /*      printf("siondefrag:                            %ld wrote left=%ld position=%" PRId64 "\n",bwrote,left,_sion_get_position(outfp)); */
          if (nullcount > MAXNULLCNT) {
            fprintf(stderr, "timeout on write data , aborting ...\n");
            exit(0);
          }
        }
      }
    }
    if (!verbose) {
      printf("\n");
    }
    sion_close(outsid);
    free(localbuffer);

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

    sion_options *options = sion_options_new();
    if (!options) {
      fprintf(stderr, "siondefrag: could not allocate options, aborting\n");
      return EXIT_FAILURE;
    }
    sion_options_set_lowlevel_api(options, SION_LOWLEVEL_API_POSIX);
    sion_options_set_keyval_mode(options, SION_KEYVAL_MODE_INLINE);
    sion_options_set_endianness(options, sion_get_file_endianness(sid) ? SION_ENDIANNESS_BIG : SION_ENDIANNESS_LITTLE);
    sion_options_set_chunksizes(options, ntasks, siondefrag_sum_bytes_per_task);
    sion_options_set_fsblksize(options, fsblksize);
    outsid = sion_open(outfilename, SION_OPEN_WRITE, ntasks, options);
    sion_options_delete(options);

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

      /* set position in outfile to current rank */
      sion_switch_logical_file(outsid, rank);

      /* seek position in infile */
      sion_switch_logical_file(sid, rank);

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

      /* loop over key-value blocks */
      while (sion_fread_key_iterator_next(sid, &key, &len) == SION_SUCCESS) {
        /* allocate buffer space */
        if ((long)localbuffer_size < (long)len) {
          if (localbuffer_size > 0) {
            free(localbuffer);
          }
          localbuffer = malloc(len);
          if (localbuffer == NULL) {
            fprintf(stderr, "cannot allocate localbuffer of size %zu , aborting ...\n", len);
            free(siondefrag_sum_bytes_per_task);
            return 1;
          }
          localbuffer_size = len;
        }

        if (verbose) {
          printf("siondefrag:                               copy now key[%12ld] with %zu bytes \n", (long)key, len);
        }

        /* read data of key */
        bread = sion_fread_key(localbuffer, key, 1, len, sid);
        /* bread=len; */
        if (bread != len) {
          fprintf(stderr, "cannot read data of key %ld with len %zu, aborting ...\n", (long)key, (size_t)len);
          free(siondefrag_sum_bytes_per_task);
          return 1;
        }

        /* write data of key */
        bwrote = sion_fwrite_key(localbuffer, key, 1, len, outsid);
        /* bread=len; */
        if (bwrote != len) {
          fprintf(stderr, "cannot write data of key %ld with len %zu, aborting ...\n", (long)key, (size_t)len);
          free(siondefrag_sum_bytes_per_task);
          return 1;
        }
      }
    }

    sion_close(outsid);
    if (!verbose) {
      printf("\n");
    }
  }

  free(siondefrag_sum_bytes_per_task);

  sion_close(sid);

  return 0;
}

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

  fprintf(stderr, "Creates a new SIONlib file <outfn> from an existing <insionfn>. The\n");
  fprintf(stderr, "resulting file has only one chunk. This can be used to make SIONlib\n");
  fprintf(stderr, "files less sparse (e.g. by using -q 1 which effectively removes file\n");
  fprintf(stderr, "system block alignment).\n\n");

  fprintf(stderr, "Options:\n");
  fprintf(stderr, "  [-Q <fsblksize>]   filessystem blocksize for new sion file in MB\n");
  fprintf(stderr, "                     (default from input file)\n");
  fprintf(stderr, "  [-q <fsblksize>]   filessystem blocksize for new sion file in bytes\n");
  fprintf(stderr, "  [-v]               verbose mode\n");
  fprintf(stderr, "  [-V]               show version of SIONlib\n");
  fprintf(stderr, "  [-h]               show this help\n");
  exit(1);
}
