/*
 * Copyright (c) 2009, Lawrence Livermore National Security, LLC.
 * Produced at the Lawrence Livermore National Laboratory.
 * Written by Adam Moody <moody20@llnl.gov>.
 * LLNL-CODE-411039.
 * All rights reserved.
 * This file is part of The Scalable Checkpoint / Restart (SCR) library.
 * For details, see https://sourceforge.net/projects/scalablecr/
 * Please also read this file: LICENSE.TXT.
*/

/* This is a utility program that lets one insert log entries into the
 * SCR log.  It's a small C program which must run on the
 * same node where rank 0 runs -- it requires the same environment as
 * the running job to identify the proper logging target.
*/

#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <string.h>
#include <errno.h>
#include <time.h>
#include <unistd.h>

#include "scr.h"
#include "scr_io.h"
#include "scr_err.h"
#include "scr_log.h"
#include "scr_param.h"

struct arglist {
  char*  username;
  char*  jobname;
  time_t start;

  int     transfer;
  char*   transfer_type;
  char*   transfer_from;
  char*   transfer_to;
  int*    transfer_ckpt;
  time_t* transfer_start;
  double* transfer_secs;
  double* transfer_bytes;
};

int scr_log_enable = 1;

int global_ckpt;
time_t global_start;
double global_secs;
double global_bytes;

void print_usage()
{
  printf("\n");
  printf("scr_log_transfer -- record a file transfer operation in the SCR log\n");
  printf("\n");
  printf("Options:\n");
  printf("  -u <username>  Username of job owner, reads $USER if not specified\n");
  printf("  -j <jobname>   Job name of job, reads $SCR_JOB_NAME if not specified\n");
  printf("  -s <seconds>   Job start time, uses current UNIX timestamp if not specified\n");
  printf("\n");
  printf("  -T <type>      Event type (string)\n");
  printf("  -X <from>      From directory (string)\n");
  printf("  -Y <to>        To directory (string)\n");
  printf("  -C <id>        Checkpoint id (integer)\n");
  printf("  -S <start>     Transfer start time as UNIX timestamp (integer)\n");
  printf("  -D <duration>  Duration in seconds (integer)\n");
  printf("  -B <bytes>     Number of bytes transfered (integer)\n");
  printf("\n");
  return;
}

int processArgs(int argc, char **argv, struct arglist* args)
{
  int i, j;
  char *argptr;
  char flag;

  /* set to default values */
  args->username = NULL;
  args->jobname  = NULL;
  args->start    = 0;

  args->transfer          = 0;
  args->transfer_type     = NULL;
  args->transfer_from     = NULL;
  args->transfer_to       = NULL;
  args->transfer_ckpt     = NULL;
  args->transfer_start    = NULL;
  args->transfer_secs     = NULL;
  args->transfer_bytes    = NULL;

  for (i=1; i<argc; i++) {
    /* check for options */
    j = 0;
    if (argv[i][j++] == '-') {
      /* flag is the first char following the '-' */
      flag = argv[i][j++];
      argptr = NULL;

      /* handles "-i#" or "-i #" */
      if (argv[i][j] != 0) {
        argptr = &(argv[i][j]);
      } else {
        argptr = argv[i+1];
        i++;
      }

      /* single argument parameters */
      if (strchr("ujsTXYCSDB", flag)) {
        switch(flag) {
        case 'u':
          args->username = strdup(argptr);
          break;
        case 'j':
          args->jobname = strdup(argptr);
          break;
        case 's':
          args->start = (time_t) strtoul(argptr, NULL, 0);
          break;

        case 'T':
          args->transfer_type = strdup(argptr);
          break;
        case 'X':
          args->transfer_from = strdup(argptr);
          break;
        case 'Y':
          args->transfer_to = strdup(argptr);
          break;
        case 'C':
          global_ckpt = (int) atoi(argptr);
          args->transfer_ckpt = &global_ckpt;
          break;
        case 'S':
          global_start = (time_t) strtoul(argptr, NULL, 0);
          args->transfer_start = &global_start;
          break;
        case 'D':
          global_secs = (double) atoi(argptr);
          args->transfer_secs = &global_secs;
          break;
        case 'B':
          global_bytes = (double) atoi(argptr);
          args->transfer_bytes = &global_bytes;
          break;
        }
        continue;
      } else {
        scr_err("Invalid flag -%c", flag);
        print_usage();
        return 0;
      }
    } else {
      scr_err("Unknown argument %s", argv[i]);
      print_usage();
      return 0;
    }
  }

  return 1;
}

int main (int argc, char *argv[])
{
  int rc = 0;

  /* process command line arguments */
  struct arglist args;
  if (!processArgs(argc, argv, &args)) {
    scr_err("scr_log_transfer: Failed to process args @ %s:%d",
            __FILE__, __LINE__
    );
    return 1;
  }

  char* value = NULL;

  /* read in job parameters (if user didn't specify them) */
  if (args.username == NULL) {
    if ((value = getenv("USER")) != NULL) {
      args.username = strdup(value);
    }
  }
  if (args.jobname == NULL) {
    if ((value = getenv("SCR_JOB_NAME")) != NULL) {
      args.jobname = strdup(value);
    }
  }
  if (args.start == 0) {
    args.start = scr_log_seconds();
  }

  /* read in log parameters */
  scr_param_init();
  if ((value = scr_param_get("SCR_LOG_ENABLE")) != NULL) {
    scr_log_enable = atoi(value);
  }
  scr_param_finalize();

  if (scr_log_enable) {
    /* init logging */
    if (scr_log_init() == SCR_SUCCESS) {
      /* register job */
      if (args.username != NULL && args.jobname != NULL) {
        if (scr_log_job(args.username, args.jobname, args.start) != SCR_SUCCESS) {
          scr_err("scr_log_transfer: Failed to register job, disabling logging @ %s:%d",
                  __FILE__, __LINE__
          );
          scr_log_enable = 0;
          rc = 1;
        }
      } else {
        scr_err("scr_log_transfer: Missing username, jobname, or start time, disabling logging @ %s:%d",
                __FILE__, __LINE__
        );
        scr_log_enable = 0;
        rc = 1;
      }
    } else {
      scr_err("scr_log_transfer: Failed to initialize SCR logging, disabling logging @ %s:%d",
              __FILE__, __LINE__
      );
      scr_log_enable = 0;
      rc = 1;
    }
  }

  if (scr_log_enable) {
  //  if (scr_log_event(args.event_type, args.event_note, args.event_ckpt, args.event_start, args.event_secs) != SCR_SUCCESS)
    if (scr_log_transfer(args.transfer_type, args.transfer_from, args.transfer_to,
          args.transfer_ckpt, args.transfer_start, args.transfer_secs, args.transfer_bytes) != SCR_SUCCESS)
    {
      rc = 1;
    }
  }

  if (scr_log_enable) {
    scr_log_finalize();
  }

  return rc;
}
