#!/usr/bin/env python3

import argparse
import datetime
import subprocess
import sys
import time

WINDOW_SIZE = 100

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.')
parser.add_argument('-i', '--interval', type=int, default=10, help='Interval in seconds between samples (default 10)')
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 followed by K/M/G/T')

args = parser.parse_args()

def parse_value(v):
    suffix = v[-1].upper()
    if suffix in SUFFIXES:
        exponent = 3*(SUFFIXES.find(suffix)+1)
        return int(v[:-1])*(10**exponent)
    else:
        return int(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

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

    current = current_val()

    samples.append(current)
    if len(samples) < 2:
        continue
    if len(samples) > WINDOW_SIZE:
        samples = samples[-WINDOW_SIZE:]

    rate = (current - samples[0])/((len(samples)-1)*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:
        print()
        break