/*
 * sc_sorter
 *
 * This is a reimplementation of sksorter.cc to use the scamper file API.
 *
 * $Id: sc_sorter.c,v 1.26 2007/05/14 04:11:02 mjl Exp $
 *
 * Copyright (C) 2006-2007 The University of Waikato
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, version 2.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 */

#if defined(__APPLE__)
#include <stdint.h>
#endif

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

#if defined(DMALLOC)
#include <dmalloc.h>
#endif

#include "scamper_tlv.h"
#include "scamper_addr.h"
#include "scamper_list.h"
#include "scamper_trace.h"
#include "scamper_ping.h"
#include "scamper_file.h"

#include "utils.h"

#include "mjl_splaytree.h"

static uint32_t               options   = 0;
static char                  *prefix    = NULL;
static char                  *infile    = NULL;
static splaytree_t           *cyclesets = NULL;
static splaytree_t           *openfiles = NULL;
static scamper_file_filter_t *filter    = NULL;

typedef struct openfile
{
  char             *filename;
  scamper_file_t   *file;
  int               active_cycles;
} openfile_t;

typedef struct cycleset
{
  scamper_cycle_t  *cycle;
  time_t            epoch;
  openfile_t      **segments;
  int               segment_cnt;
} cycleset_t;

#define OPT_PREFIX  0x00000001 /* P: */
#define OPT_HELP    0x00000002 /* ?: */

static void usage(const char *argv0, uint32_t opt_mask)
{
  fprintf(stderr, "usage: %s [-?] [-F prefix] <infile>\n", argv0);

  if(opt_mask == 0) return;

  fprintf(stderr, "\n");

  if(opt_mask & OPT_HELP)
    fprintf(stderr, "    -? give an overview of the usage of sc_sorter\n");

  if(opt_mask & OPT_PREFIX)
    fprintf(stderr, "    -F root path prefix to use\n");

  return;
}

static int check_options(int argc, char *argv[])
{
  char    ch;
  size_t  len;
  char   *opts = "F:?";
  char   *opt_prefix = NULL;

  while((ch = getopt(argc, argv, opts)) != -1)
    {
      switch(ch)
	{
	case 'F':
	  options |= OPT_PREFIX;
	  opt_prefix = optarg;
	  break;

	case '?':
	default:
	  usage(argv[0], 0xffffffff);
	  return -1;
	}
    }

  /*
   * if a prefix has been specified, then ensure that the string has a
   * trailing forward slash
   */
  if(options & OPT_PREFIX)
    {
      if((len = strlen(opt_prefix)) == 0)
	{
	  usage(argv[0], OPT_PREFIX);
	  return -1;
	}

      if(opt_prefix[len-1] != '/')
	{
	  if((prefix = malloc(len + 2)) == NULL)
	    {
	      return -1;
	    }
	  memcpy(prefix, opt_prefix, len);
	  prefix[len] = '/';
	  prefix[len+1] = '\0';
	}
      else
	{
	  if((prefix = strdup(opt_prefix)) == NULL)
	    {
	      return -1;
	    }
	}
    }

  /*
   * ensure that an input file is specified
   */
  if(argc - optind != 1)
    {
      usage(argv[0], 0);
      return -1;
    }

  infile = argv[optind];

  return 0;
}

/*
 * create_filename
 *
 * this function creates a filename (and full path to it) given a cycle and
 * segment value.  the filename is created as follows:
 *
 * <PREFIX>/<LIST>/<MONITOR>/YYYY/MM/<MONITOR>.<LIST>.YYYYMMDD_<SEGMENT>
 */
static char *create_filename(char *filename, size_t len,
			     scamper_cycle_t *cycle, int segment)
{
  struct tm *tm;
  time_t tt;

  tt = cycle->start_time;
  tm = gmtime(&tt);

  snprintf(filename, len, "%s%s/%s/%d/%02d/%s.%03d.%d%02d%02d_%03d",
	   prefix != NULL ? prefix : "",
	   cycle->list->name,
	   cycle->list->monitor,
	   tm->tm_year + 1900,
	   tm->tm_mon + 1,
	   cycle->list->monitor,
	   cycle->list->id,
	   tm->tm_year + 1900,
	   tm->tm_mon + 1,
	   tm->tm_mday,
	   segment);

  return filename;
}

/*
 * openfile_alloc
 *
 * allocate a record that allows us to associate a filename with an open
 * scamper file, which can then be shared as necessary.
 *
 * caller is responsible for putting the structure into splaytree.
 */
static openfile_t *openfile_alloc(char *filename)
{
  openfile_t *openfile = NULL;

  if((openfile = malloc_zero(sizeof(openfile_t))) == NULL)
    {
      goto err;
    }

  if((openfile->filename = strdup(filename)) == NULL)
    {
      goto err;
    }

  return openfile;

 err:
  if(openfile != NULL)
    {
      if(openfile->file != NULL) scamper_file_close(openfile->file);
      free(openfile);
    }
  return NULL;
}

/*
 * openfile_close
 *
 *
 */
static void openfile_close(void *vopenfile)
{
  openfile_t *openfile = (openfile_t *)vopenfile;
  char a[1024], b[1024];
  struct stat sb;
  mode_t mode;

  /* close the file */
  if(openfile->file != NULL)
    {
      scamper_file_close(openfile->file);
      openfile->file = NULL;
    }

  if(openfile->active_cycles == 0)
    {
      snprintf(a, sizeof(a), "%s.warts", openfile->filename);
      snprintf(b, sizeof(b), "%s.inprogress", openfile->filename);

      /* make sure there is no .warts file there currently */
      if(stat(a, &sb) != 0 && errno == ENOENT)
	{
	  rename(b, a);

	  mode = S_IRUSR | S_IRGRP | S_IROTH;
	  chmod(a, mode);
	}
    }

  if(openfile->filename != NULL) free(openfile->filename);
  free(openfile);

  return;
}

static int openfile_cmp(const void *va, const void *vb)
{
  return strcmp(((const openfile_t *)va)->filename,
		((const openfile_t *)vb)->filename);
}

/*
 * file_open
 *
 * open a warts file for appending, using the .inprogress extension.  if
 * an .inprogress file does not exist, but a .warts one does, then rename
 * the .warts file to .inprogress and open that.
 */
static scamper_file_t *file_open(char *filename)
{
  struct stat sb;
  char a[1024], b[1024];
  int rv;

  /* if the .inprogress file exists, use that */
  snprintf(a, sizeof(a), "%s.inprogress", filename);
  if(stat(a, &sb) == 0)
    {
      return scamper_file_open(a, 'a', "warts");
    }
  else if(errno != ENOENT)
    {
      return NULL;
    }

  /*
   * if the .inprogress file does not exist, check if one ending in .warts
   * does.  if one does exist, then rename it to .inprogress and open it.
   */
  snprintf(b, sizeof(b), "%s.warts", filename);
  if((rv = stat(b, &sb)) == 0 || errno == ENOENT)
    {
      if(rv == 0 && (rename(b, a) != 0 || chmod(a, S_IRUSR | S_IWUSR) != 0))
	{
	  return NULL;
	}

      return scamper_file_open(a, 'a', "warts");
    }

  return NULL;
}

/*
 * openfile_get
 *
 * return an openfile structure corresponding to the segment number for
 * this cycle's cycleset.
 */
static openfile_t *openfile_get(cycleset_t *cycleset, time_t sec)
{
  openfile_t **segments, findme, *openfile = NULL;
  int segment;
  int segment_cnt;
  char *delim, filename[1024];

  /*
   * if the timestamp is before the epoch, just write the record into
   * segment zero.
   */
  if(sec <= cycleset->epoch)
    {
      segment = 0;
    }
  else
    {
      segment = (sec - cycleset->epoch) / 86400;
    }

  /* if we need to increase the size of the segment table, then do so */
  if(cycleset->segment_cnt <= (segment_cnt = segment+1))
    {
      segments = (openfile_t **)realloc(cycleset->segments,
					segment_cnt * sizeof(openfile_t *));
      if(segments == NULL)
	{
	  return NULL;
	}
      cycleset->segments = segments;

      memset(cycleset->segments + cycleset->segment_cnt,
	     0, (segment_cnt - cycleset->segment_cnt) * sizeof(openfile_t *));

      cycleset->segment_cnt = segment_cnt;
    }

  /*
   * if there is no open file recorded for this segment, then try and
   * acquire one
   */
  if(cycleset->segments[segment] == NULL)
    {
      /* this is the full path and name of file we'll be opening */
      create_filename(filename, sizeof(filename), cycleset->cycle, segment);

      /* see if we already have the file open */
      findme.filename = filename;
      if((openfile = splaytree_find(openfiles, &findme)) != NULL)
	{
	  cycleset->segments[segment] = openfile;
	  return openfile;
	}

      /*
       * get a pointer to the last / in the filename, and ensure the
       * full directory path has been created.
       */
      delim = string_lastof_char(filename, '/');
      *delim = '\0';
      if(mkdir_wrap(filename, 0755) == -1)
	{
	  goto err;
	}
      *delim = '/';

      /*
       * allocate an 'openfile' structure that wraps the filename and points
       * to an open scamper_file for it
       */
      if((openfile = openfile_alloc(filename)) == NULL)
	{
	  goto err;
	}

      /* save the record of this open file in the splaytree */
      if(splaytree_insert(openfiles, openfile) == NULL)
	{
	  goto err;
	}

      /* now actually try and open the file */
      if((openfile->file = file_open(filename)) == NULL)
	{
	  perror("could not open file");
	  goto err;
	}

      cycleset->segments[segment] = openfile;
    }

  return cycleset->segments[segment];

 err:
  if(openfile != NULL) openfile_close(openfile);
  return NULL;
}

/*
 * cycleset_alloc
 *
 * allocate a cycleset structure, and determine the 24-hour epoch for the
 * cycle
 */
static cycleset_t *cycleset_alloc(scamper_cycle_t *cycle)
{
  cycleset_t *cycleset;
  struct tm tm;
  time_t tt;

  if((cycleset = malloc_zero(sizeof(cycleset_t))) == NULL)
    {
      return NULL;
    }

  cycleset->cycle = scamper_cycle_use(cycle);

  /* determine the 24-hour epoch for the cycle */
  tt = cycle->start_time;
  memcpy(&tm, gmtime(&tt), sizeof(tm));

  cycleset->epoch = tt - (3600 * tm.tm_hour) - (60 * tm.tm_min) - tm.tm_sec;

  return cycleset;
}

static int cycleset_cmp(const void *va, const void *vb)
{
  const cycleset_t *a = (const cycleset_t *)va;
  const cycleset_t *b = (const cycleset_t *)vb;

  if(a->cycle < b->cycle) return -1;
  if(a->cycle > b->cycle) return  1;

  return 0;  
}

static void cycleset_free(cycleset_t *cycleset)
{
  if(cycleset == NULL)
    {
      return;
    }

  /* free our reference of the cycle */
  if(cycleset->cycle != NULL)
    {
      scamper_cycle_free(cycleset->cycle);
    }

  /* free any open files related to this cycleset */
  if(cycleset->segments != NULL)
    {
      free(cycleset->segments);
    }

  free(cycleset);

  return;
}

static cycleset_t *cycleset_get(scamper_cycle_t *cycle)
{
  cycleset_t findme, *cycleset;

  findme.cycle = cycle;

  if((cycleset = splaytree_find(cyclesets, &findme)) == NULL)
    {
      if((cycleset = cycleset_alloc(cycle)) == NULL)
	{
	  return NULL;
	}

      if(splaytree_insert(cyclesets, cycleset) == NULL)
	{
	  cycleset_free(cycleset);
	  return NULL;
	}
    }

  return cycleset;
}

/*
 * cycle_openfile_get
 *
 * given a cycle and a time since the cycle start point, return an openfile
 * record.
 */
static openfile_t *cycle_openfile_get(scamper_cycle_t *cycle, time_t sec)
{
  cycleset_t *cycleset;

  /* first, get the corresponding cycleset */
  if((cycleset = cycleset_get(cycle)) == NULL)
    {
      return NULL;
    }

  /* second, get the actual output file for this segment */
  return openfile_get(cycleset, sec);
}

/*
 * handle_trace
 *
 * write the trace to the appropriate output file.
 */
static int handle_trace(scamper_trace_t *trace)
{
  openfile_t *openfile;
  int rc = 0;

  /* skip over traces without a cycle attached */
  if(trace->cycle == NULL)
    {
      goto done;
    }

  if((openfile = cycle_openfile_get(trace->cycle,trace->start.tv_sec)) == NULL)
    {
      rc = -1; goto done;
    }
	  
  if(scamper_file_write_trace(openfile->file, trace) == -1)
    {
      rc = -1; goto done;
    }

 done:
  scamper_trace_free(trace);
  return rc;
}

/*
 * handle_ping
 *
 * write the ping object to the appropriate output file.
 */
static int handle_ping(scamper_ping_t *ping)
{
  openfile_t *openfile;
  int rc = 0;

  /* skip over traces without a cycle attached */
  if(ping->cycle == NULL)
    {
      goto done;
    }

  if((openfile = cycle_openfile_get(ping->cycle, ping->start.tv_sec)) == NULL)
    {
      rc = -1; goto done;
    }
	  
  if(scamper_file_write_ping(openfile->file, ping) == -1)
    {
      rc = -1; goto done;
    }

 done:
  scamper_ping_free(ping);
  return rc;
}

/*
 * handle_cycle_start
 *
 * write the start record to the output file.  since this is a newly active
 * cycle, increment the active cycles parameter for the openfile.
 */
static int handle_cycle_start(scamper_cycle_t *cycle)
{
  openfile_t *openfile;
  int rc = 0;

  if((openfile = cycle_openfile_get(cycle, cycle->start_time)) == NULL)
    {
      rc = -1; goto done;
    }

  if(scamper_file_write_cycle_start(openfile->file, cycle) == -1)
    {
      rc = -1; goto done;
    }

  openfile->active_cycles++;

 done:
  scamper_cycle_free(cycle);
  return rc;
}

/*
 * handle_cycle_stop
 *
 * a cycle has concluded.  write the cycle_stop record, and reduce the
 * active-cycles count.
 */
static int handle_cycle_stop(scamper_cycle_t *cycle)
{
  openfile_t *openfile;
  int rc = 0;

  if((openfile = cycle_openfile_get(cycle, cycle->stop_time)) == NULL)
    {
      rc = -1; goto done;
    }

  if(scamper_file_write_cycle_stop(openfile->file, cycle) == -1)
    {
      rc = -1; goto done;
    }

  openfile->active_cycles--;

 done:
  scamper_cycle_free(cycle);
  return rc;
}

/*
 * handle_cycle_def
 *
 * a cycle definition in an input file probably means that a cycle was
 * defined, but that the current input file was first used in the middle
 * of the cycle.  the only thing sc_sorter has to do is increment the
 * number of known active cycles for the openfile.
 */
static int handle_cycle_def(scamper_cycle_t *cycle)
{
  openfile_t *openfile;
  int rc = 0;

  if((openfile = cycle_openfile_get(cycle, cycle->start_time)) == NULL)
    {
      rc = -1; goto done;
    }

  openfile->active_cycles++;

 done:
  scamper_cycle_free(cycle);
  return rc;
}

static void cleanup(void)
{
  if(prefix != NULL)
    {
      free(prefix);
      prefix = NULL;
    }

  if(filter != NULL)
    {
      scamper_file_filter_free(filter);
      filter = NULL;
    }

  if(cyclesets != NULL)
    {
      splaytree_free(cyclesets, (splaytree_free_t)cycleset_free);
      cyclesets = NULL;
    }

  if(openfiles != NULL)
    {
      splaytree_free(openfiles, NULL);
      openfiles = NULL;
    }

  return;
}

int main(int argc, char *argv[])
{
  uint16_t filter_types[] = {
    SCAMPER_FILE_OBJ_TRACE,
    SCAMPER_FILE_OBJ_PING,
    SCAMPER_FILE_OBJ_CYCLE_START,
    SCAMPER_FILE_OBJ_CYCLE_STOP,
    SCAMPER_FILE_OBJ_CYCLE_DEF,
  };
  uint16_t filter_cnt = sizeof(filter_types)/sizeof(uint16_t);
  scamper_file_t *file;
  void *data;
  uint16_t type;

#if defined(DMALLOC)
  free(malloc(1));
#endif

  atexit(cleanup);

  if(check_options(argc, argv) == -1)
    {
      return -1;
    }

  if((cyclesets = splaytree_alloc(cycleset_cmp)) == NULL)
    {
      return -1;
    }

  if((openfiles = splaytree_alloc(openfile_cmp)) == NULL)
    {
      return -1;
    }

  if((file = scamper_file_open(infile, 'r', NULL)) == NULL)
    {
      usage(argv[0], 0);
      fprintf(stderr, "could not open infile %s\n", infile);
      return -1;
    }

  if((filter = scamper_file_filter_alloc(filter_types, filter_cnt)) == NULL)
    {
      fprintf(stderr, "could not allocate filter\n");
      return -1;
    }

  while(scamper_file_read(file, filter, &type, &data) == 0)
    {
      if(data == NULL)
	{
	  /* EOF */
	  break;
	}

      if(type == SCAMPER_FILE_OBJ_TRACE)
	{
	  if(handle_trace(data) != 0)
	    {
	      return -1;
	    }
	}
      else if(type == SCAMPER_FILE_OBJ_PING)
	{
	  if(handle_ping(data) != 0)
	    {
	      return -1;
	    }
	}
      else if(type == SCAMPER_FILE_OBJ_CYCLE_START)
	{
	  if(handle_cycle_start(data) != 0)
	    {
	      return -1;
	    }
	}
      else if(type == SCAMPER_FILE_OBJ_CYCLE_STOP)
	{
	  if(handle_cycle_stop(data) != 0)
	    {
	      return -1;
	    }
	}
      else if(type == SCAMPER_FILE_OBJ_CYCLE_DEF)
	{
	  if(handle_cycle_def(data) != 0)
	    {
	      return -1;
	    }
	}
    }

  scamper_file_close(file);

  /* close all open files */
  splaytree_free(openfiles, openfile_close);
  openfiles = NULL;

  return 0;
}