#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Plugin installation:
# wget "https://pasternok.org/sources/nut_batt_runtime.py?raw" -O /etc/munin/plugins/nut_batt_runtime
# chmod +x /etc/munin/plugins/nut_batt_uptime
# echo -e "\n[nut_batt_runtime]\ngroup nut" >> /etc/munin/plugin-conf.d/munin-node
# systemctl restart munin-node

### config
UPSC_BIN = "/bin/upsc"
UPSMON_CONFIG = "/etc/nut/upsmon.conf"
RUNTIME_UNIT = "m"  # h - hours, m - minutes, s - seconds

### code - do not change anything below here
__author__ = "Damian Pasternok <my_forename at pasternok.org>"
__version__ = "0.2.0"
__date__ = "2022/02/10"
__copyright__ = "Copyright (C) 2022 Damian Pasternok"
__license__ = "GPL"

import argparse
import io
import re
import shlex
from subprocess import Popen, PIPE, DEVNULL


class Upsc(object):
    def __init__(self):
        upses = self.list_upses()
        self.upsc = {}

        for ups in upses:
            raw_data = io.StringIO(
                Popen(
                    shlex.split(f"{UPSC_BIN} {ups}"),
                    stdout=PIPE,
                    stderr=DEVNULL,
                    universal_newlines=True,
                )
                .communicate()[0]
                .strip()
            )
            self.upsc[ups] = {}

            for param in raw_data:
                l_param = param.rstrip().split(":")
                self.upsc[ups][l_param[0]] = l_param[1].lstrip()

            if not self.upsc[ups].get("battery.runtime"):
                del self.upsc[ups]

    def list_upses(self):
        upses = []
        with open(UPSMON_CONFIG, "r") as fh:
            for l in fh:
                if re.search("^MONITOR", l):
                    upses.append(l.split()[1])
        return upses


class Munin(Upsc):
    def __init__(self):
        super(Munin, self).__init__()

        if RUNTIME_UNIT == "h":
            self.rt_label = "hours"
        elif RUNTIME_UNIT == "m":
            self.rt_label = "minutes"
        elif RUNTIME_UNIT == "s":
            self.rt_label = "seconds"
        else:
            self.rt_label = "s"

    def format_rt_value(self, value, unit):
        if unit == "h":
            return int(value) / 3600
        elif unit == "m":
            return int(value) / 60
        elif unit == "s":
            return value
        else:
            return value

    def config(self):
        out = f"""graph_title UPS battery runtime
graph_args --logarithmic --units=si
graph_scale no
graph_category ups
graph_vlabel {self.rt_label}"""
        for i, ups in enumerate(self.upsc):
            out += f"""
runtime{i}.label {self.upsc[ups].get("ups.model") or f"Generic UPS #{i}"}"""

        return out

    def data(self):
        out = ""
        for i, ups in enumerate(self.upsc):
            out += f"""runtime{i}.value {self.format_rt_value(self.upsc[ups]["battery.runtime"], RUNTIME_UNIT)}\n"""

        return out.rstrip()


def main():
    munin = Munin()

    parser = argparse.ArgumentParser(description="NUT Battery Runtime Munin Plugin.")
    parser.add_argument("config", help="print graph configuration", nargs="?")
    args = parser.parse_args()

    if args.config == "config":
        print(munin.config())
    elif not args.config:
        print(munin.data())
    else:
        parser.print_help()


if __name__ == "__main__":
    main()