#!/usr/bin/python3

import os
import sys
import subprocess
import glob
import fnmatch

class NoSize(Exception):
    pass

def kmgt(sz, frac=1):
    t = {
        'K': pow(1024, 1),
        'M': pow(1024, 2),
        'G': pow(1024, 3),
        'T': pow(1024, 4),
        '': 1}

    for k in sorted(t, key=t.__getitem__, reverse=True):
        if sz >= t[k]:
            n = sz / float(t[k])
            tpl = "{:." + str(frac) + "f}{}"
            return tpl.format(n, k)

def my_du(path, skip):
    total = 0

    if os.path.isfile(path):
        return os.stat(path).st_size

    for root, dirs, files in os.walk(path, topdown=False):
        for f in files:
            filepath = os.path.join(root, f)
            if should_skip(filepath, skip):
                continue
            try:
                if os.path.isfile(filepath) and not os.path.islink(filepath):
                    total += os.path.getsize(filepath)
            except OSError as e:
                pass
    return total

def sys_du(path, skip):
    cmdline = ['du', '-bs', path]
    for sg in skip.split(' '):
        cmdline.append('--exclude={}'.format(sg))

    print('# run: {}'.format(cmdline))
    r = subprocess.run(cmdline, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if r.returncode not in [0, 1]:
        # exit code 1 is OK, when file disappeared during process
        sys.stderr.write("cmd: {} returncode: {}".format(cmdline, r.returncode))
        sys.stderr.write(r.stderr)
        return -1

    if not r.stdout:
        raise NoSize

    sz = int(r.stdout.decode().split('\t')[0])
    return sz

def du(path, skip):
    if use_du:
        return sys_du(path, skip)
    else:
        return my_du(path, skip)

def should_skip(path, skip):

    for g in skip.split(' '):
        if fnmatch.fnmatch(path, g):
            return True
    return False


# read parameters
prefix = os.getenv('PREFIX')
basename = os.getenv('BASENAME')

pathlist = os.getenv('PATHLIST', '/root /home/*')
maxlim = os.getenv('MAXLIM', '1G')
skip = os.getenv('SKIP', '')
fullpath = int(os.getenv('FULLPATH', '0'))
use_du = int(os.getenv('USE_DU', '1'))

for path in pathlist.split(' '):
    for curpath in glob.glob(path):
        if should_skip(curpath, skip):
            print('# skip {}'.format(curpath))
            continue
        print('# count {}'.format(curpath))

        try:
            sz = du(curpath, skip)
        except NoSize as e:
            print('# Skip {} (NoSize)'.format(curpath))
            continue

        print("NAME: {}{}".format(prefix, curpath if fullpath else os.path.basename(curpath)))
        print("TAGS: dirsize")
        print("METHOD: numerical|maxlim={}|minlim=0".format(maxlim))
        print("DETAILS: {}".format(kmgt(sz)))
        print("STATUS: {}".format(sz))
        print()
