#!/usr/bin/env python3

import argparse
import datetime
import math
import subprocess
import sys
import time

SUFFIXES = 'KMGT'

parser = argparse.ArgumentParser(description='Display progress and time estimates for arbitrary tasks')
parser.add_argument('-f', '--final', type=str, help='Expected final size/value at completion, optionally with a K/M/G/T suffix')
parser.add_argument('-i', '--interval', type=int, default=10, help='Interval in seconds between samples (default 10)')
parser.add_argument('-w', '--window', type=int, default=100, help='Number of samples to keep for the sliding window (default 100)')
parser.add_argument('-d', '--decay', type=float, default=1, help='Decay coefficient for older samples (default 1). Must be between 0 and 1 inclusive. The lower this is, the more responsive/swingy the estimate will be.')
parser.add_argument('-V', '--termination-value-threshold', type=float, default=0.95, metavar='FRACTION', help='Fraction of the expected final value that must be reached in order to terminate (default 0.95). Reaching this threshold is necessary but not sufficient, see also -I')
parser.add_argument('-I', '--termination-inactivity-threshold', type=int, default=10, metavar='COUNT', help='Number of consecutive unchanged samples that must be observed in order to terminate (default 10). Reaching this threshold is necessary but not sufficient, see also -V')
parser.add_argument('-l', '--log-file', type=str, metavar='PATH', help='File to log the time series to. Will be saved as a csv, with columns for timestamp and for value. Will append data if the file already exists.')
tracker_types = parser.add_mutually_exclusive_group(required=True)
tracker_types.add_argument('-p', '--path', type=str, help='Track total disk usage of a given path')
tracker_types.add_argument('-c', '--command', type=str, help='Track value returned by a shell command; this should return a single number, optionally with a K/M/G/T suffix')

args = parser.parse_args()
if args.interval < 1:
    print('interval must be at least 1', file=sys.stderr)
    exit(1)
if args.window < 1:
    print('window size must be at least 1', file=sys.stderr)
    exit(1)
if args.decay < 0 or args.decay > 1:
    print('decay coefficient must be between 0 and 1', file=sys.stderr)
    exit(1)
if args.termination_value_threshold < 0 or args.termination_value_threshold > 1:
    print('termination value threshold must be between 0 and 1', file=sys.stderr)
    exit(1)
if args.termination_inactivity_threshold < 0:
    print('termination inactivity threshold must be nonnegative', file=sys.stderr)
    exit(1)

def parse_value(v):
    suffix = v[-1].upper()
    if suffix in SUFFIXES:
        exponent = 3*(SUFFIXES.find(suffix)+1)
        return float(v[:-1])*(10**exponent)
    else:
        return float(v)

def display_value(v):
    suffix = ''
    for c in SUFFIXES:
        if v < 10**3:
            break
        v = v / 10**3
        suffix = c
    return '{:.1f}{}'.format(v, suffix)

def display_timedelta(d):
    result = ''
    if d.days != 0:
        result += '{}d'.format(d.days)
    result += '{}h'.format(d.seconds // 3600)
    result += '{:02}m'.format((d.seconds % 3600) // 60)
    return result

if args.path:
    def current_val_helper():
        du = subprocess.run(['du', '--bytes', '--summarize', args.path], capture_output=True, text=True).stdout
        return parse_value(du.split()[0])
else:
    def current_val_helper():
        result = subprocess.run(args.command, shell=True, capture_output=True, text=True).stdout
        return parse_value(result.strip())

if args.log_file:
    log_file = open(args.log_file, mode='a', buffering=1)
else:
    log_file = None

def current_val():
    result = current_val_helper()
    if log_file:
        print('{},{}'.format(datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S"), result), file=log_file)
    return result

if args.final:
    final = parse_value(args.final)
    if final <= 0:
        print('final value must be positive', file=sys.stderr)
        exit(1)
else:
    final = None

deltas = []
current = current_val()
while True:
    time.sleep(args.interval)

    new = current_val()

    deltas.append(new - current)
    current = new
    if len(deltas) > args.window:
        deltas = deltas[-args.window:]

    total = 0
    divisor = 0
    coeff = 1
    for d in deltas[::-1]:
        total += coeff*d
        divisor += coeff
        coeff *= args.decay
    rate = total/(divisor*args.interval)

    print('\033[2K\r', end='')
    print('{} - {}/s'.format(display_value(current), display_value(rate)), end='')
    if final:
        fraction = current / final
        value_remaining = final - current
        print(' - {} total - {:.1f}% complete'.format(display_value(final), 100*fraction), end='')
        seconds_remaining = (value_remaining / rate) if rate > 0 else math.inf
        if seconds_remaining < 0:
            print(' - unknown time remaining', end='')
        elif seconds_remaining > 864000:
            print(' - stalled', end='')
        else:
            time_remaining = datetime.timedelta(seconds=seconds_remaining)
            eta = datetime.datetime.now() + time_remaining
            print(' - {} remaining - ETA {}'.format(
                display_timedelta(time_remaining),
                eta.isoformat(sep=' ', timespec='minutes'),
            ), end='')
    sys.stdout.flush()

    if (
        final
        and current >= args.termination_value_threshold*final
        and len(deltas) >= args.termination_inactivity_threshold
        and all(d == 0 for d in deltas[-args.termination_inactivity_threshold:])
    ):
        print()
        break