import collectd,os,subprocess,sys,re,time

# https://www.programcreek.com/python/example/106897/collectd.register_read

PATHS = [] # { name, dir, exclude }
INTERVAL = 60 * 60 * 12 # twice a day. change values in docker-graphite if you change this

def du(pathinfo):
    cmd = 'timeout 1800 du -Dsb "{}"'.format(pathinfo['dir'])
    if pathinfo['exclude'] != '':
        cmd += ' --exclude "{}"'.format(pathinfo['exclude'])

    collectd.info('computing size with command: %s' % cmd);
    try:
        size = subprocess.check_output(cmd, shell=True).split()[0].decode('utf-8')
        collectd.info('\tsize of %s is %s (time: %i)' % (pathinfo['dir'], size, int(time.time())))
        return size
    except Exception as e:
        collectd.info('\terror getting the size of %s: %s' % (pathinfo['dir'], str(e)))
        return 0

def parseSize(size):
    units = {"B": 1, "KB": 10**3, "MB": 10**6, "GB": 10**9, "TB": 10**12}
    number, unit, _ = re.split('([a-zA-Z]+)', size.upper())
    return int(float(number)*units[unit])

def dockerSize():
    # use --format '{{json .}}' to dump the string. '{{if eq .Type "Images"}}{{.Size}}{{end}}' still creates newlines
    cmd = 'timeout 1800 docker system df --format "{{.Size}}" | head -n1'
    try:
        size = subprocess.check_output(cmd, shell=True).strip().decode('utf-8')
        collectd.info('size of docker images is %s (%s) (time: %i)' % (size, parseSize(size), int(time.time())))
        return parseSize(size)
    except Exception as e:
        collectd.info('error getting docker images size : %s' % str(e))
        return 0

# configure is called for each module block. this is called before init
def configure(config):
    global PATHS

    for child in config.children:
        if child.key != 'Path':
            collectd.info('du plugin: Unknown config key "%s"' % key)
            continue

        pathinfo = { 'name': '', 'dir': '', 'exclude': '' }
        for node in child.children:
            if node.key == 'Instance':
                pathinfo['name'] = node.values[0]
            elif node.key == 'Dir':
                pathinfo['dir'] = node.values[0]
            elif node.key == 'Exclude':
                pathinfo['exclude'] = node.values[0]

        PATHS.append(pathinfo);
        collectd.info('du plugin: monitoring %s' % pathinfo['dir']);

def init():
    global PATHS
    collectd.info('custom du plugin initialized with %s %s' % (PATHS, sys.version))

def read():
    for pathinfo in PATHS:
        size = du(pathinfo)

        # type comes from https://github.com/collectd/collectd/blob/master/src/types.db
        val = collectd.Values(type='capacity', plugin='du', plugin_instance=pathinfo['name'])
        val.dispatch(values=[size], type_instance='usage')

    size = dockerSize()
    val = collectd.Values(type='capacity', plugin='du', plugin_instance='docker')
    val.dispatch(values=[size], type_instance='usage')



collectd.register_init(init)
collectd.register_config(configure)
collectd.register_read(read, INTERVAL)