#!/usr/bin/env python3

import argparse
import datetime
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.')
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)

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():
        du = subprocess.run(['du', '--bytes', '--summarize', args.path], capture_output=True, text=True).stdout
        return parse_value(du.split()[0])
else:
    def current_val():
        result = subprocess.run(args.command, shell=True, capture_output=True, text=True).stdout
        return parse_value(result.strip())

if args.final:
    final = parse_value(args.final)
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='')
        if rate > 0:
            time_remaining = datetime.timedelta(seconds=(value_remaining / rate))
            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 >= final and all(d == 0 for d in deltas[-10:]):
        print()
        break