/*
  pipemeter - A program to show status of a pipe

  Copyright Clint Byrum 2002, All Rights Reserved.

  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; either version 2 of the License, or
  (at your option) any later version.

  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

  ---

*/

#include <sys/time.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>

#include <stdio.h>
#include <signal.h>
#include <sys/time.h>
#include <errno.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <limits.h>
#include <time.h>

#define VERSION "0.7.2"

#define DEFAULT_BLOCK_SIZE 8192
#define DEFCOLS 70
//#define PBAR "[-----------------------------------------------------]"
#define PBCHAR '-'
#define PBLEFT '['
#define PBRIGHT ']'
#define PBFILL '*'
//#define STARS "*****************************************************"
//#define PBAR_LEN 53
// This must be changed if averages and rate formats are changed
#define OTHER_LEN 23
// Might want to turn down for slower machines.
// This defines the number of samples to average.
#define LAST_MAX 20

// TODO: async/multi threaded to seperate reading and writing.

void show_just_rate(int sig);
void show_progress(int sig);
void parseopts(int argc, char *argv[],char **buffer);
void formatbytes(char *obuffer,double b);
size_t full_write (int desc, const char *ptr, size_t len);
double avg_bytes(long long bytes);

// XXX: hmm... should we be using size_t instead?
// XXX: ok.. probably too many globals .. may have to clean this up

extern char **environ;

long long bytes;
long long lastbytes = 0;
long long filesize = 0;
double itimer_seconds;
long block_size;
char *filename;
int report_mode=0;
//time_t start_time;
struct timeval start_time;

char *progressbar;
char *progressfill;
unsigned int pbarlen;

// Used for calculating a more pertinent average.
long long last[LAST_MAX];
unsigned char lcnt; // only used for indexing the above array

int main(int argc, char *argv[]) {
  struct itimerval it;

  //FIXME: these should be initialized! (sig_empty or something?)
  // (I have a book I don't feel like looking for right now that explains it)
  struct sigaction sa,sa_orig;

  struct timeval interval;
  long bytesin,bytesout;
  char *buffer;
  int in,out;
  unsigned int columns,i;
  char *colstr;
  out=fileno(stdout);
  lastbytes=0;
  bytes=0;

  parseopts(argc,argv,&buffer);
  if(filename == NULL) {
    in=fileno(stdin);
  } else {
    //fprintf(stderr, "opening %s for reading\n\n\n", filename);
    in = open(filename, O_RDONLY);
    if(in < 0) {
      perror("can't open input file");
      exit(2);
    }
    if(filesize) {
      fprintf(stderr, "Warning: -s overrides size of file given by -f!\n");
    } else {
      struct stat s;
      if(fstat(in, &s) != 0) {
        perror("fstat");
        exit(3);
      }
      filesize = s.st_size;
      //fprintf(stderr, "filesize is %lld\n", filesize);
      /* note: trying to use fstat on a file descriptor opened to /dev/zero
         results in a st_size of 0, which by coincidence disables the
         progress bar. This is exactly what I wanted, but I don't know
         how portable it is.

         Also: giving -f file on the command line shows the progress bar for
         regular files, with no way to disable it. This is *not* what I
         wanted, but will be OK for now.  -- B.
      */
    }
  }

  // XXX: hrm.. why not just make columns a long.. ? :P
  errno=0;
  colstr=getenv("COLUMNS");
  if(colstr) {
    columns=(unsigned int)strtol(getenv("COLUMNS"),NULL,10);
  } else {
    columns=DEFCOLS;
  }
  if(errno) {
    //fprintf(stderr,"Invalid number of columns specified, setting to 70");
    columns=DEFCOLS;
  }
  pbarlen=columns-OTHER_LEN;
  progressbar=(char *)malloc(pbarlen*sizeof(char)+1);
  //strcpy(progressbar,PBAR);
  progressbar[0]=PBLEFT;
  // yes I'm aware it would be faster to set a limit before the loop. This is
  // only ever going to factor in if somebody has a 30000 column display
  for(i=1;i<pbarlen-3;i++) {
    progressbar[i]=PBCHAR;
  }
  progressbar[i]=PBRIGHT;
  progressbar[i+1]='\0';

  progressfill=(char *)malloc(pbarlen*sizeof(char));
  for(i=0;i<pbarlen-4;i++) {
    progressfill[i]=PBFILL;
  }
  progressfill[i]=PBRIGHT;
  progressfill[i+1]='\0';
//fprintf(stderr,"Bleh %d [%d] %d",i,pbarlen,strlen(progressfill));
  
  sa.sa_handler= filesize ? show_progress : show_just_rate;
  sigaction(SIGTERM, &sa, &sa_orig);
  sigaction(SIGINT,  &sa, &sa_orig);
  if(!report_mode) {
    // setup timers
    sigaction(SIGALRM, &sa, &sa_orig);
  
    interval.tv_usec = itimer_seconds * 1000000;
    interval.tv_sec  = 0;

    it.it_value=interval;
    it.it_interval=interval;
    if(setitimer(ITIMER_REAL,&it,NULL)) {
      printf("error setting itimer %d\n",errno);
      exit(1);
    }
  }

  // Loop until we get an EOF on stdin
  //start_time=time(NULL);
  gettimeofday(&start_time,NULL);
  while((bytesin=read(in,buffer,block_size))) {
    if(bytesin > 0) {
      bytes += bytesin;
      bytesout=full_write(out,buffer,bytesin);
    } else {
      /*
       * If there's no data available, we get an EINTR error. This is not bad.
       * This is, quite possibly, Linux specific. Need to try other platforms.
      */
#ifdef EINTR
      if(errno!=EINTR) {
        fprintf(stderr,"ERROR: read error [%d], aborting.\n",errno);
        exit(1);
      }
#else
      fprintf(stderr,"ERROR: read error [%d], aborting.\n",errno);
      exit(1);
#endif
    }
  }

  // Final status report
  if(filesize) {
    show_progress(0);
  } else {
    show_just_rate(0);
  }
  fprintf(stderr, "\n");
  return 0;
}

void show_just_rate(int sig) {
  /* if no progress bar, we use this function */
  char numbuf[512]; // XXX: hrm... bad form, but so much simpler
  double bytespersecond;
  static unsigned char tcount=0;
  struct timeval tnow;
  double tx,ty;
  if(sig != SIGALRM) {
    // If we're reporting
    gettimeofday(&tnow,NULL);
    tx=tnow.tv_sec+(tnow.tv_usec * 0.000001);
    ty=start_time.tv_sec+(start_time.tv_usec * 0.000001);
    itimer_seconds = tx - ty;
    lastbytes=0.0;
  }
  if(tcount < LAST_MAX) {
    bytespersecond=(double)(bytes-lastbytes)/(double)itimer_seconds;
    avg_bytes(bytes-lastbytes);
    tcount++;
  } else {
    bytespersecond=avg_bytes(bytes-lastbytes);
  }
  formatbytes(numbuf,bytespersecond);
  fprintf(stderr,"%s/s",numbuf);
  /* Show total bytes through */
  /* Thanks to Sean Reifschneider for this idea */
  formatbytes(numbuf,bytes);
  fprintf(stderr," %s",numbuf);
  lastbytes=bytes;
  if(sig == 0 || sig==SIGTERM || sig==SIGINT) {
    fprintf(stderr,"\n");
    exit(0);
  } else {
    fprintf(stderr,"\r");
  }
}

void show_progress(int sig) {
  //char buf[512]; // XXX: yes, this is bad form.
  char buf2[512]; // XXX: hopefully we fix them before there are 100 or so
  
  double percent  = (double)bytes / (double)filesize * 100;
  int    progress;
  double bytespersecond ;
  
  struct timeval tnow;
  double tx,ty;

  progress = (double)bytes / (double)filesize * pbarlen;
  if(sig != SIGALRM) {
    // If we're reporting
    gettimeofday(&tnow,NULL);
    tx=tnow.tv_sec+(tnow.tv_usec * 0.000001);
    ty=start_time.tv_sec+(start_time.tv_usec * 0.000001);
    itimer_seconds = tx - ty;
    lastbytes=0.0;
  }

  bytespersecond = (double)(bytes-lastbytes)/(double)itimer_seconds;
  formatbytes(buf2,bytespersecond);
  lastbytes=bytes;

  if(progress > pbarlen) {
    // fprintf(stderr, "Ow! -s filesize parameter (%lld) overflowed!\n", filesize);
    progress = pbarlen; // just keep printing full progress bars and warnings.
  }

  strncpy(progressbar+1,progressfill,progress);
  fprintf(stderr, progressbar);
  fprintf(stderr," %s/s",buf2);
  formatbytes(buf2,bytes);
  fprintf(stderr," %s",buf2);
  fprintf(stderr, " %2.1f%%\r", percent);
  if(sig == 0 || sig==SIGTERM || sig==SIGINT) {
    fprintf(stderr,"\n");
    exit(0);
  } else {
    fprintf(stderr,"\r");
  }
}

void parseopts(int argc, char *argv[], char **buffer) {
  int c;
  // Set some defaults
  itimer_seconds=1;
  block_size=DEFAULT_BLOCK_SIZE;
  filename = NULL;
  do {
    int mult = 1;
    c=getopt(argc,argv,"f:s:b:i:rV");
    switch(c) {
    case -1:
      // No more options
      break;
    case 0:
      // No options passed, this is fine.
      break;
    case 'b':
      // blocksize qualifiers by Ian McMahon, tmbg@hardcoders.org, 9/07/2002
      // we're gonna examine the last char, and if it matches [kKmMgG] we'll act accordingly
      switch(optarg[strlen(optarg) - 1]) { 
        case 'g': /* FALLTHRU */
        case 'G': /* FALLTHRU */
          mult *= 1024;  
        case 'm': /* FALLTHRU */
        case 'M': /* FALLTHRU */
          mult *= 1024;  
        case 'k': /* FALLTHRU */
        case 'K': 
          mult *= 1024;  
          optarg[strlen(optarg) - 1] = '\0';
          break;
        default:
          break;
      }

      block_size=strtol(optarg,NULL,10);
      block_size *= mult;

      if(block_size==LONG_MIN || block_size==LONG_MAX || block_size <= 0) {
        fprintf(stderr,"Bad block size: %s\n",optarg);
        exit(1);
      }
      break;
    case 'i':
      itimer_seconds=strtod(optarg,NULL);
      if(itimer_seconds <= 0.0) {
        fprintf(stderr,"Bad interval: %s\n",optarg);
        exit(1);
      }
      break;
    case 's':
      // Copied from the 'b' option - Clint Byrum 10/10/2002
      switch(optarg[strlen(optarg) - 1]) { 
        case 'g': /* FALLTHRU */
        case 'G': /* FALLTHRU */
          mult *= 1024;  
        case 'm': /* FALLTHRU */
        case 'M': /* FALLTHRU */
          mult *= 1024;  
        case 'k': /* FALLTHRU */
        case 'K': 
          mult *= 1024;  
          optarg[strlen(optarg) - 1] = '\0';
          break;
        default:
          break;
      }
      
      filesize=strtol(optarg,NULL,10);
      filesize *= mult;
      
      if(filesize == LONG_MAX || filesize == LONG_MIN || filesize <= 0) {
        fprintf(stderr, "Bad filesize: %s\n", optarg);
        exit(1);
      }
      break;
    case 'f':
      filename = optarg;
      break;
    case 'r':
      report_mode=1;
      break;
    case 'V':
      // Exiting here because otherwise we'll start trying to read/write
      fprintf(stderr,"pipemeter v%s\n",VERSION);
      exit(0);
    case '?':
    case ':':
      // XXX: better errors
    default:
      fprintf(stderr,"usage: pipemeter { -b blocksize } { -i interval }\n");
      //fprintf(stderr,"debug: %c %d\n",(char)c,c);
      exit(1);
    }
  } while(c > 0);
  *buffer=(char *)malloc(block_size*sizeof(char));
}

void formatbytes(char *obuffer,double b) {
  double tmp;
  if(b>1073741824) {
    tmp=b/1073741824;
    // If you can get it to go faster than 999.99G/s ... You win.
    sprintf(obuffer,"%6.2fG",tmp);
  } else if(b>1048576) {
    tmp=b/1048576;
    sprintf(obuffer,"%6.2fM",tmp);
  } else if(b>2048) { // under 2k, let it be
    tmp=b/1024;
    sprintf(obuffer,"%6.2fk",tmp);
  } else
    sprintf(obuffer,"%6.2fB",b);
}

/* This is used to get a smoother average */
// TODO: moving average?
double avg_bytes(long long bytes) {
  long long tmp=0;
  unsigned char i;
  last[lcnt]=bytes;
  if(lcnt == LAST_MAX) {
    lcnt=0;
  } else {
    lcnt++;
  }
  for(i=0;i<LAST_MAX;i++) {
    tmp += last[i];
  }
  return (((double)tmp/(double)LAST_MAX)*(1.0/itimer_seconds));
}

/* taken from fileutils... yay GPL (and thanks GNU for the code) */
size_t
full_write (int desc, const char *ptr, size_t len)
{
  size_t total_written = 0;

  while (len > 0)
    {
      ssize_t written = write (desc, ptr, len);
      if (written <= 0)
	{
	  /* Some buggy drivers return 0 when you fall off a device's end.  */
	  if (written == 0)
	    errno = ENOSPC;
#ifdef EINTR
	  if (errno == EINTR)
	    continue;
#endif
	  break;
	}
      total_written += written;
      ptr += written;
      len -= written;
    }
  return total_written;
}