import collectd,os,subprocess,json,re

# https://blog.dbrgn.ch/2017/3/10/write-a-collectd-python-plugin/

def parseSiSize(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 parseBinarySize(size):
    units = {"B": 1, "KIB": 2**10, "MIB": 2**20, "GIB": 2**30, "TIB": 2**40}
    number, unit, _ = re.split('([a-zA-Z]+)', size.upper())
    return int(float(number)*units[unit])

def init():
    collectd.info('custom docker-status plugin initialized')

def read():
    try:
        lines = subprocess.check_output('docker stats --format  "{{ json . }}" --no-stream --no-trunc', shell=True).decode('utf-8').strip().split("\n")
    except Exception as e:
        collectd.info('\terror getting docker stats: %s' % (str(e)))
        return 0

    # Sample line
    # {"BlockIO":"430kB / 676kB","CPUPerc":"0.00%","Container":"7eae5e6f4f11","ID":"7eae5e6f4f11","MemPerc":"59.15%","MemUsage":"45.55MiB / 77MiB","Name":"1062eef3-ec96-4d81-9f02-15b7dd81ccb9","NetIO":"1.5MB / 3.48MB","PIDs":"5"}

    for line in lines:
        stat = json.loads(line)
        containerName = stat["Name"] # same as app id

        # currently we only collect data for apps main containers. Those have the app id as the Name which is 36 long
        if len(containerName) != 36:
            continue

        networkData = stat["NetIO"].split("/")
        networkRead = parseSiSize(networkData[0].strip())
        networkWrite = parseSiSize(networkData[1].strip())

        blockData = stat["BlockIO"].split("/")
        blockRead = parseSiSize(blockData[0].strip())
        blockWrite = parseSiSize(blockData[1].strip())

        memUsageData = stat["MemUsage"].split("/")
        memUsed = parseBinarySize(memUsageData[0].strip())
        memMax = parseBinarySize(memUsageData[1].strip())

        cpuPercData = stat["CPUPerc"].strip("%")
        cpuPerc = float(cpuPercData)

        # type comes from https://github.com/collectd/collectd/blob/master/src/types.db and https://collectd.org/wiki/index.php/Data_source
        val = collectd.Values(type='gauge', plugin='docker-stats', plugin_instance=containerName)
        val.dispatch(values=[networkRead], type_instance='network-read')
        val.dispatch(values=[networkWrite], type_instance='network-write')
        val.dispatch(values=[blockRead], type_instance='blockio-read')
        val.dispatch(values=[blockWrite], type_instance='blockio-write')
        val.dispatch(values=[memUsed], type_instance='mem-used')
        val.dispatch(values=[memMax], type_instance='mem-max')
        val.dispatch(values=[cpuPerc], type_instance='cpu-perc')

        val = collectd.Values(type='counter', plugin='docker-stats', plugin_instance=containerName)
        val.dispatch(values=[networkRead], type_instance='network-read')
        val.dispatch(values=[networkWrite], type_instance='network-write')
        val.dispatch(values=[blockRead], type_instance='blockio-read')
        val.dispatch(values=[blockWrite], type_instance='blockio-write')

collectd.register_init(init)
# see Interval setting in collectd.conf for polling interval
collectd.register_read(read)