#!/usr/bin/env python
#
# This should be both valid Python 2 and 3, as the collectd Python
# plugin can be either Python 2 or 3 (decided at compile time).
#
# Oddly, the Arch 'collectd' package chose Python 2, despite using
# Python 3 as the default system Python.  I assume they have a good
# reason for it, though I can't find it discussed anywhere.

import sys
import dbus

# For compatibility with Python 2 & 3
if sys.version_info.major is 2:
    def iteritems(d):
        return d.iteritems()
else:
    def iteritems(d):
        return d.items()

# DBus utility function
def get_property(dbus_iface, propname):
    return dbus_iface.proxy_object.Get(dbus_iface.dbus_interface, propname, dbus_interface='org.freedesktop.DBus.Properties')

# The main program

def init():
    global system_manager
    system_manager = dbus.Interface(dbus.SystemBus().get_object("org.freedesktop.systemd1", "/org/freedesktop/systemd1"), "org.freedesktop.systemd1.Manager")

def read():
    global system_manager

    dat = {
        # These are the values of the
        # basic/unit-name.h:UnitActiveStates enum.
        'NActiveUnits': 0,
        'NReloadingUnits': 0,
        'NInactiveUnits': 0,
        'NFailedUnits': 0,
        'NActivatingUnits': 0,
        'NDeactivatingUnits': 0
    }
    for unit in system_manager.ListUnits():
        key = str('N%sUnits' % unit[3].capitalize())
        dat[key] = dat[key] + 1

    for prop in ['NJobs', 'NInstalledJobs', 'NFailedJobs']:
        dat[prop] = int(get_property(system_manager, prop))

    return dat

if __name__ == '__main__':
    init()
    for (key, val) in iteritems(read()):
        print("%s => %s" % (key, val))
else:
    import collectd
    def _config(conf):
        pass
    def _read():
        for (key, val) in iteritems(read()):
            collectd.info("---- dump() ----")
            collectd.Values(
                plugin='systemd',
                type='count',
                type_instance=key
            ).dispatch(
                values=[val]
            )
    collectd.register_config(_config)
    collectd.register_init(init)
    collectd.register_read(_read)