#!/usr/bin/env python3
# This file is part of Checkmk (https://checkmk.com). It is subject to the terms and
# conditions defined in the file COPYING, which is part of this source code package.

# NOTE: Careful when replacing the *-import below with a more specific import. This can cause
# problems because it might remove variables from the check-context which are necessary for
# resolving legacy discovery results such as [("SUMMARY", "diskstat_default_levels")]. Furthermore,
# it might also remove variables needed for accessing discovery rulesets.
import json
from cmk.base.check_api import LegacyCheckDefinition
from cmk.base.config import check_info
from cmk.agent_based.v2 import Service

#<<<graylog_input_metrics:sep(0)>>>
# {"641e88d05d447a677efde199": {"input_state": "FAILED", "input_name": "kafka_cef_test",
# "input_type": "CEF Kafka", "input_port": null, "im_m1_rate": 0.0, "im_m5_rate": 0.0,
# "im_m15_rate": 0.0, "rs_m1_rate": 0.0, "rs_m5_rate": 0.0, "rs_m15_rate": 0.0},
# "641e32885d447a677efd2dbf": {"input_state": "RUNNING", "input_name": "UDP-test",
# "input_type": "Syslog UDP", "input_port": 1514, "im_m1_rate": 1.0846244336700077,
# "im_m5_rate": 1.3700826278955827, "im_m15_rate": 1.254406787430692, "rs_m1_rate": 145.45579305762527,
# "rs_m5_rate": 180.6486220431909, "rs_m15_rate": 165.26666376319292},
# "641e32795d447a677efd2d9e": {"input_state": "RUNNING", "input_name": "testTCP", "input_type": "Syslog TCP",
# "input_port": 1515, "im_m1_rate": 1.057872514816615, "im_m5_rate": 1.364957693749168,
# "im_m15_rate": 1.2528742858546844, "rs_m1_rate": 140.4719944116262, "rs_m5_rate": 178.57816158901215,
# "rs_m15_rate": 163.80530659055356}}

def parse_graylog_input_metrics(section):
    parsed = json.loads(section[0][0])
    return parsed


def inventory_graylog_input_metrics(parsed):
    for input_id, input_info in parsed.items():
        input_name = input_info["input_name"]
        yield Service(item=f"{input_name} ({input_id})")


def check_graylog_input_metrics(item, params, parsed):
    item_id = item.split()[-1][1:-1]
    input_info = parsed.get(item_id)
    if not input_info:
        return

    input_state = input_info["input_state"]

    state = 1
    if input_state == "RUNNING": state = 0
    elif input_state == "FAILED": state = 2

    yield state, "State: %s" % input_state
    yield 0, "Type: %s" % input_info["input_type"]

    if input_info["input_port"]:
        yield 0, "Port: %s" % input_info["input_port"]

    for key, dsname, unit, unit_func, infotext in [
        ("im_m1_rate",  "im_m1_rate",  "/1min",  float, "Incoming messages"),
        ("im_m5_rate",  "im_m5_rate",  "/5min",  float, ""),
        ("im_m15_rate", "im_m15_rate", "/15min", float, ""),
        ("rs_m1_rate",  "rs_m1_rate",  "/1min",  get_bytes_human_readable, "Incoming data"),
        ("rs_m5_rate",  "rs_m5_rate",  "/5min",  get_bytes_human_readable, ""),
        ("rs_m15_rate", "rs_m15_rate", "/15min", get_bytes_human_readable, ""),
    ]:
        value = input_info[key]
        value = round(value, 2)

        rate_upper = params.get(key, {}).get("%s_upper" % key, (None, None))
        rate_lower = params.get(key, {}).get("%s_lower" % key, (None, None))

        yield check_levels(
            value,
            dsname,
            rate_upper + rate_lower,
            human_readable_func=unit_func,
            unit=unit,
            infoname=infotext
        )


# A customer wanted us to support <= for "below", not <. This meant copying
# check_levels() and child functions wholesale out of
# lib/python3/cmk/base/check_api.py just to change < to <= in this function.
# Sometimes a sledgehammer is what it takes... D:
# Start copy ===================================================================
from cmk.agent_based import v1 as _v1

def _do_check_levels(value, levels, human_readable_func):
    warn_upper, crit_upper, warn_lower, crit_lower = levels

    # Critical cases
    if crit_upper is not None and value >= crit_upper:
        return 2, _levelsinfo_ty("at", warn_upper, crit_upper, human_readable_func)
    if crit_lower is not None and value <= crit_lower:
        return 2, _levelsinfo_ty("below", warn_lower, crit_lower, human_readable_func)

    # Warning cases
    if warn_upper is not None and value >= warn_upper:
        return 1, _levelsinfo_ty("at", warn_upper, crit_upper, human_readable_func)
    if warn_lower is not None and value <= warn_lower:
        return 1, _levelsinfo_ty("below", warn_lower, crit_lower, human_readable_func)
    return 0, ""

def _levelsinfo_ty(ty, warn, crit, human_readable_func):
    warn_str = "never" if warn is None else f"{human_readable_func(warn)}"
    crit_str = "never" if crit is None else f"{human_readable_func(crit)}"
    return f" (warn/crit {ty} {warn_str}/{crit_str})"

def _build_perfdata(dsname, value, levels, boundaries):
    used_boundaries = boundaries if isinstance(boundaries, tuple) and len(boundaries) == 2 else ()
    return [(dsname, value, levels[0], levels[1], *used_boundaries)]

def check_levels(value, dsname, params, unit, human_readable_func, infoname, boundaries=None):
    def render_func(x):
        return "%s%s" % (human_readable_func(x), unit)

    if params and isinstance(params, dict):
        result, *metrics = _v1.check_levels_predictive(
            value,
            levels=params,
            metric_name=dsname,
            render_func=render_func,
            label=infoname,
            boundaries=boundaries,
        )
        assert isinstance(result, _v1.Result)
        return (
            int(result.state),
            result.summary,
            [
                (m.name, m.value, *m.levels, *m.boundaries)
                for m in metrics
                if isinstance(m, _v1.Metric)
            ],
        )

    infotext = f"{render_func(value)}"
    if infoname:
        infotext = f"{infoname}: {infotext}"

    levels = params

    state, levelstext = _do_check_levels(value, levels, render_func)
    state, levelstext = _do_check_levels(value, levels, render_func)
    return (
        state,
        infotext + levelstext,
        _build_perfdata(dsname, value, levels, boundaries),
    )
# End copy ====================================================================


check_info["graylog_input_metrics"] =  LegacyCheckDefinition(
    parse_function = parse_graylog_input_metrics,
    check_function = check_graylog_input_metrics,
    discovery_function = inventory_graylog_input_metrics,
    service_name = "Graylog Input %s",
    check_ruleset_name = "graylog_input_metrics",
    check_default_parameters={},
)
