#!/usr/bin/env python3

"""prometheus-liquidctl-exporter – host a metrics HTTP endpoint with Prometheus formatted data from liquidctl

This is an experimental script that collects stats from liquidctl and exposes them as a http://localhost:8098/metrics
endpoint in the Prometheus text format.
See: https://prometheus.io/docs/instrumenting/exposition_formats/#text-format-example

Example metric with labels:
# HELP liquidctl liquidctl exported metrics
# TYPE liquidctl gauge
liquidctl{device="NZXT Kraken X (X42, X52, X62 or X72)",sensor="liquid_temperature",unit="°C"} 33.6

Usage:
  prometheus-liquidctl-exporter [options]

Options:
  --legacy-690lc                     Use Asetek 690LC in legacy mode (old Krakens)
  --server-port <number>             Port for the HTTP /metrics endpoint
  -v, --verbose                      Output additional information
  -g, --debug                        Show debug information on stderr
  --version                          Display the version number
  --help                             Show this message

Device selection options (see: list -v):
  -m, --match <substring>            Filter devices by description substring
  -n, --pick <number>                Pick among many results for a given filter
  --vendor <id>                      Filter devices by hexadecimal vendor ID
  --product <id>                     Filter devices by hexadecimal product ID
  --release <number>                 Filter devices by hexadecimal release number
  --serial <number>                  Filter devices by serial number
  --bus <bus>                        Filter devices by bus
  --address <address>                Filter devices by address in bus
  --usb-port <port>                  Filter devices by USB port in bus


Copyright Alex Berryman, Jonas Malaco and contributors
SPDX-License-Identifier: GPL-3.0-or-later
"""

import logging
import sys
import time
import usb
from datetime import timedelta

from docopt import docopt
from liquidctl.driver import *
from prometheus_client import start_http_server
from prometheus_client.core import GaugeMetricFamily, REGISTRY, InfoMetricFamily

LOGGER = logging.getLogger(__name__)


def gauge_name_sanitize(name):
    return name.replace(" ", "_").lower()


class LiquidCollector(object):
    def __init__(self):
        self.description = "liquidctl exported metrics"

    def collect(self):
        labels = ["device", "sensor", "unit"]
        g = GaugeMetricFamily("liquidctl", self.description, labels=labels)
        i = InfoMetricFamily("liquidctl", self.description, labels=["device"])
        for d in devs:
            try:
                get_status = d.get_status()
                for metric in get_status:
                    sanitized_name = gauge_name_sanitize(metric[0])
                    sample_value = metric[1]
                    unit = metric[2]

                    if isinstance(sample_value, timedelta):
                        # cast timedelta into seconds and override the supplied unit
                        sample_value = sample_value.seconds
                        unit = "seconds"

                    if unit != "":
                        # FIXME doesn't handle multiple equal devices well
                        label_values = [d.description, sanitized_name, unit]
                        g.add_metric(label_values, value=sample_value)
                        LOGGER.debug(
                            "%s: %s as GaugeMetric %s labels %s",
                            d.description,
                            metric,
                            sanitized_name,
                            "/".join(label_values),
                        )
                    else:
                        i.add_metric([d.description], value={sanitized_name: sample_value})
                        LOGGER.debug(
                            "%s: %s InfoMetric labeled with %s => %s",
                            d.description,
                            metric,
                            sanitized_name,
                            sample_value,
                        )
            except usb.core.USBError as err:
                LOGGER.warning("failed to read from the device, possibly serving stale data")
                LOGGER.debug(err, exc_info=True)
        yield g
        yield i


def _make_opts(arguments):
    options = {}
    for arg, val in arguments.items():
        if val is not None and arg in _PARSE_ARG:
            opt = arg.replace("--", "").replace("-", "_")
            options[opt] = _PARSE_ARG[arg](val)
    return options


_PARSE_ARG = {
    "--legacy-690lc": bool,
    "--vendor": lambda x: int(x, 16),
    "--product": lambda x: int(x, 16),
    "--release": lambda x: int(x, 16),
    "--serial": str,
    "--bus": str,
    "--address": str,
    "--usb-port": lambda x: tuple(map(int, x.split("."))),
    "--match": str,
    "--pick": int,
}

if __name__ == "__main__":
    args = docopt(__doc__, version="0.1.1")
    opts = _make_opts(args)
    devs = list(find_liquidctl_devices(**opts))
    for d in devs:
        LOGGER.info("initializing %s", d.description)
        d.connect()

    if args["--debug"]:
        args["--verbose"] = True
        logging.basicConfig(level=logging.DEBUG, format="[%(levelname)s] %(name)s: %(message)s")
    elif args["--verbose"]:
        logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
    else:
        logging.basicConfig(level=logging.WARNING, format="%(message)s")
        sys.tracebacklimit = 0

    REGISTRY.register(LiquidCollector())

    if args["--server-port"]:
        server_port = int(args["--server-port"])
    else:
        server_port = 8098

    start_http_server(server_port)
    LOGGER.debug("server started on port %s", server_port)

    try:
        while True:
            # Keep HTTP server alive in a loop
            time.sleep(2)
    except KeyboardInterrupt:
        LOGGER.info("canceled by user")
    finally:
        for d in devs:
            try:
                LOGGER.info("disconnecting from %s", d.description)
                d.disconnect()
            except:
                LOGGER.exception("unexpected error when disconnecting")
