/* progress -- copy stdin to stdout while displaying a progress meter on stderr
 *
 * To do: Could be made a little more efficient by not checking the
 * style inside display(), but having separate loops for each
 * combination of options.
 *
 * To do: Error handling for errors when reading stdin or writing
 * stdout.
 *
 * Created: 29 Nov 2009
 * Author: Bert Bos <bert@w3.org>
 */

#define _GNU_SOURCE 1		/* We need LLONG_MAX and SA_RESTART */
/* #include "config.h" -- Not currently used */
#include <limits.h>
#include <stdlib.h>
#include <sys/time.h>
#include <time.h>
#include <sys/errno.h>
#include <unistd.h>
#include <stdio.h>
#include <err.h>
#include <sysexits.h>
#include <assert.h>
#include <signal.h>
#include <string.h>

#ifndef LLONG_MAX
#define LLONG_MAX LONG_LONG_MAX	/* Some versions of GNU C define this */
#endif

#define OPTIONS "[-h] [-s style] [-m max] [-i interval]"
#define KILO 1024LL
#define MEGA (1024 * KILO)
#define GIGA (1024 * MEGA)
#define TERA (1024 * GIGA)

typedef enum _Valtype { BYTE, LINE, SECOND } Valtype;
typedef enum _Style { DOT, SPIN, SIZE, LINES, PERCENT, BAR, SPEED } Style;

/* Global variables, used by alarm_handler() and main() */
static long long int nbytes = 0, nlines = 0, starttime, max = 0;
static Valtype maxtype = BYTE;
static Style style = SIZE;


/* help -- print help text */
static void help(char *progname)
{
  printf("\
  Copy input to output and display a progress meter while doing so. Usage:\n\
    %1$s %2$s\n\
  -h\n\
    This help text.\n\
  -s style\n\
    Type of progress meter: dot (a growing row of dots), spin (a rotating\n\
    dash), size (the number of copied bytes), lines (the number of copied\n\
    lines), percent (the percentage copied, requires -m), bar (the\n\
    percentage and a growing bar), speed (number of bytes per second).\n\
    Default: size.\n\
  -m max\n\
    Expected total number of bytes, lines or seconds. Add s for\n\
    seconds, l for lines.\n\
  -i interval\n\
    How often to print something. Add s for seconds, l for lines.\n\
    Default: 0.5s\n\
  Numbers may be followed by K, M, G or T (or k, g, m, t) to multiply them\n\
  by 1024, etc. E.g., to print a percentage of 30 MiB every 2 seconds:\n\
    %1$s -s percent -m 30M -i 2s\n", progname, OPTIONS);
}


/* usage -- print usage message and exit */
static void usage(char *progname)
{
  fprintf(stderr, "Usage: %s %s  (Try -h for help)\n", progname, OPTIONS);
  exit(EX_USAGE);
}

/* runtime -- return # of microseconds since starttime */
static long long int runtime(long long int starttime)
{
  struct timeval curtime;

  (void) gettimeofday(&curtime, NULL);
  return 1000000 * curtime.tv_sec + curtime.tv_usec - starttime;
}


/* display -- display progress in given style */
static void display(Style style, long long int nbytes, long long int nlines,
		    long long int starttime, long long int max, Valtype maxtype)
{
  static char *spinner = "-\\|/";
  static int i = 0;
  double p;
  int j;

  switch (style) {
  case DOT:
    putc('.', stderr);
    break;
  case SPIN:
    putc(spinner[i], stderr);
    putc('\b', stderr);
    i = (i + 1) % 4;
    break;
  case SIZE:
    if (nbytes > TERA)
      fprintf(stderr, "%7.2f TiB\b\b\b\b\b\b\b\b\b\b\b", 1.0 * nbytes/TERA);
    else if (nbytes > GIGA)
      fprintf(stderr, "%7.2f GiB\b\b\b\b\b\b\b\b\b\b\b", 1.0 * nbytes/GIGA);
    else if (nbytes > MEGA)
      fprintf(stderr, "%7.2f MiB\b\b\b\b\b\b\b\b\b\b\b", 1.0 * nbytes/MEGA);
    else if (nbytes > KILO)
      fprintf(stderr, "%7.2f KiB\b\b\b\b\b\b\b\b\b\b\b", 1.0 * nbytes/KILO);
    else
      fprintf(stderr, "%7lld B\b\b\b\b\b\b\b\b\b", nbytes);
    /* Bug: fails when nbytes >= 10000 * TERA */
    break;
  case LINES:
    fprintf(stderr, "%10lld lines\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b", nlines);
    /* Bug: fails when nlines >= 1,000,000,000 */
    break;
  case PERCENT:
    switch (maxtype) {
    case BYTE: fprintf(stderr, "%3.0f%%\b\b\b\b", 100.0 * nbytes / max); break;
    case LINE: fprintf(stderr, "%3.0f%%\b\b\b\b", 100.0 * nlines / max); break;
    case SECOND: fprintf(stderr,"%3.0f%%\b\b\b\b",100.0*runtime(starttime)/max);
    }
    break;
  case BAR:
    p = 100.0 * (maxtype == BYTE ? nbytes :
		 maxtype == LINE ? nlines :
		 runtime(starttime)) / max;
    fprintf(stderr, "\r%3.0f%% |", p);
    for (j = 0; j < p/2 + 0.5 && j < 50; j++) putc('=', stderr);
    if (j < 50) {putc('>', stderr); j++;}
    for (; j < 50; j++) putc('-', stderr);
    fprintf(stderr, "| %c", spinner[i]);
    i = (i + 1) % 4;
    break;
  case SPEED:
    p = 1000000.0 * nbytes / runtime(starttime);
    if (p > TERA)
      fprintf(stderr, "%7.2f TiB/s\b\b\b\b\b\b\b\b\b\b\b\b\b", p / TERA);
    else if (p > GIGA)
      fprintf(stderr, "%7.2f GiB/s\b\b\b\b\b\b\b\b\b\b\b\b\b", p / GIGA);
    else if (p > MEGA)
      fprintf(stderr, "%7.2f MiB/s\b\b\b\b\b\b\b\b\b\b\b\b\b", p / MEGA);
    else if (p > KILO)
      fprintf(stderr, "%7.2f KiB/s\b\b\b\b\b\b\b\b\b\b\b\b\b", p / KILO);
    else
      fprintf(stderr, "%7.2f B/s\b\b\b\b\b\b\b\b\b\b\b", p);
  }
}


/* alarm_handler -- signal handler function */
static void alarm_handler(int sig)
{
  /* Not sure if it is safe to do I/O here, but it avoids making the
     main loop non-blocking... At least we know that the main loop
     will not create output itself. */
  display(style, nbytes, nlines, starttime, max, maxtype);
}


/* start_timer -- start an interval alarm */
static void start_timer(long long int microseconds)
{
  struct itimerval timer;
  struct sigaction sigact;

  sigact.sa_handler = alarm_handler;
  sigact.sa_flags = SA_RESTART;
  sigemptyset(&sigact.sa_mask);
  if (sigaction(SIGALRM, &sigact, NULL) != 0) err(EX_OSERR, "");
  timer.it_interval.tv_sec = timer.it_value.tv_sec = microseconds / 1000000;
  timer.it_interval.tv_usec = timer.it_value.tv_usec = microseconds % 1000000;
  setitimer(ITIMER_REAL, &timer, NULL);
}


/* stop_timer -- stop the alarm */
static void stop_timer(void)
{
  struct itimerval timer;

  signal(SIGALRM, SIG_IGN);
  timer.it_interval.tv_sec = timer.it_value.tv_sec = 0;
  timer.it_interval.tv_usec = timer.it_value.tv_usec = 0;
  setitimer(ITIMER_REAL, &timer, NULL);
}


/* parse_style -- parse a string as a display style */
static void parse_style(const char *s, Style *style)
{
  if (strcmp(s, "dot") == 0) *style = DOT;
  else if (strcmp(s, "spin") == 0) *style = SPIN;
  else if (strcmp(s, "size") == 0) *style = SIZE;
  else if (strcmp(s, "lines") == 0) *style = LINES;
  else if (strcmp(s, "percent") == 0) *style = PERCENT;
  else if (strcmp(s, "bar") == 0) *style = BAR;
  else if (strcmp(s, "speed") == 0) *style = SPEED;
  else errx(EX_USAGE, "Unknown style (\"%s\"). Try -h for help.", s);
}


/* parse_val -- parse a string as a value and type */
static void parse_val(const char *s, long long int *val, Valtype *tp)
{
  char *p;
  double v;

  errno = 0;
  v = strtod(s, &p);
  if (v <= 0 || errno) errx(EX_USAGE, "Invalid number. Try -h for help.");
  switch (*p) {
  case 'k': case 'K': v *= KILO; p++; break;
  case 'm': case 'M': v *= MEGA; p++; break;
  case 'g': case 'G': v *= GIGA; p++; break;
  case 't': case 'T': v *= TERA; p++; break;
  }
  /* Get the type. If seconds, convert to microseconds */
  switch (*p) {
  case 'l': case 'L': *tp = LINE; p++; break;
  case 's': case 'S': *tp = SECOND; p++; v *= 1000000; break;
  default: *tp = BYTE; break;
  }
  if (*p) errx(EX_USAGE, "Unrecognized value (\"%s\"). Try -h for help.", s);
  if (v > LLONG_MAX) errx(EX_USAGE, "Number too large. Try -h for help.");
  /* Convert to long long int */
  *val = v;
}


/* main -- main body */
int main(int argc, char *argv[])
{
  Valtype steptype = SECOND;
  struct timeval curtime;
  long long int step = 500000;	/* = 0.5 second */
  int c;

  while ((c = getopt(argc, argv, "hs:m:i:")) != -1) {
    switch (c) {
    case 'h': help(argv[0]); exit(0); break;
    case 's': parse_style(optarg, &style); break;
    case 'm': parse_val(optarg, &max, &maxtype); break;
    case 'i': parse_val(optarg, &step, &steptype); break;
    default: usage(argv[0]);
    }
  }
  /* Check command line arguments */
  if (optind < argc) usage(argv[0]);
  if ((style == PERCENT || style == BAR) && max == 0)
    errx(EX_USAGE, "Option -m is required for this style. Try -h for help.");

  if (gettimeofday(&curtime, NULL) != 0) err(EX_OSERR, "Cannot get the time.");
  starttime = 1000000 * curtime.tv_sec + curtime.tv_usec;

  if (steptype == SECOND) start_timer(step);

  /* Copy input to output, displaying progress at regular intervals */
  while ((c = getchar()) != EOF) {
    nbytes++;
    if (c == '\n') nlines++;
    if ((steptype == BYTE && nbytes % step == 0) ||
	(steptype == LINE && nlines % step == 0))
      display(style, nbytes, nlines, starttime, max, maxtype);
    putchar(c);
  }
  if (steptype == SECOND) stop_timer();
  if (style != DOT) display(style, nbytes, nlines, starttime, max, maxtype);
  putc('\n', stderr);
  return 0;
}