#!/usr/bin/env python3

import itertools
import json

from cmk.agent_based.v2 import (
    AgentSection,
    CheckPlugin,
    Metric,
    Result,
    Service,
    State,
    check_levels,
    render,
)


# ---------------------------------------------------------------------------
# JSON parsing
# ---------------------------------------------------------------------------

def parse_p1(string_table):
    flat = list(itertools.chain.from_iterable(string_table))
    raw = "".join(flat)

    try:
        data = json.loads(raw)
    except json.JSONDecodeError:
        return None

    if not isinstance(data, dict):
        return None

    return data


agent_section_p1 = AgentSection(
    name="p1",
    parse_function=parse_p1,
)


# ---------------------------------------------------------------------------
# P1 tariff
# ---------------------------------------------------------------------------

def discover_p1_tariff(section):
    if section is not None and "active_tariff" in section:
        yield Service()


def check_p1_tariff(section):
    tariff = section.get("active_tariff")

    if tariff is None:
        yield Result(
            state=State.UNKNOWN,
            summary="Active tariff is not available",
        )
        return

    yield Result(
        state=State.OK,
        summary=f"Tariff: {tariff}",
    )


check_plugin_p1_tariff = CheckPlugin(
    name="p1_tariff",
    sections=["p1"],
    service_name="P1 Tariff",
    discovery_function=discover_p1_tariff,
    check_function=check_p1_tariff,
)


# ---------------------------------------------------------------------------
# Power
# ---------------------------------------------------------------------------

POWER_ITEMS = {
    "Total": "active_power_w",
    "L1": "active_power_l1_w",
    "L2": "active_power_l2_w",
    "L3": "active_power_l3_w",
}


def discover_p1_power(section):
    for item, key in POWER_ITEMS.items():
        if key in section:
            yield Service(item=item)


def check_p1_power(item, section):
    key = POWER_ITEMS.get(item)

    if key is None or key not in section:
        yield Result(
            state=State.UNKNOWN,
            summary="Power value is not available",
        )
        return

    value = float(section[key])

    yield Result(
        state=State.OK,
        summary=f"{value:.0f} W",
    )

    yield Metric(
        name="p1_power_w",
        value=value,
        boundaries=(0, None),
    )


check_plugin_p1_power = CheckPlugin(
    name="p1_power",
    sections=["p1"],
    service_name="P1 Power %s",
    discovery_function=discover_p1_power,
    check_function=check_p1_power,
)


# ---------------------------------------------------------------------------
# Voltage
# ---------------------------------------------------------------------------

VOLTAGE_ITEMS = {
    "L1": "active_voltage_l1_v",
    "L2": "active_voltage_l2_v",
    "L3": "active_voltage_l3_v",
}


def discover_p1_voltage(section):
    for item, key in VOLTAGE_ITEMS.items():
        if key in section:
            yield Service(item=item)


def check_p1_voltage(item, params, section):
    key = VOLTAGE_ITEMS.get(item)

    if key is None or key not in section:
        yield Result(
            state=State.UNKNOWN,
            summary="Voltage value is not available",
        )
        return

    value = float(section[key])

    yield from check_levels(
        value,
        levels_upper=params["levels"],
        metric_name="p1_voltage_v",
        label="Voltage",
        render_func=lambda v: f"{v:.1f} V",
        boundaries=(0, 300),
    )


check_plugin_p1_voltage = CheckPlugin(
    name="p1_voltage",
    sections=["p1"],
    service_name="P1 Voltage %s",
    discovery_function=discover_p1_voltage,
    check_function=check_p1_voltage,
    check_default_parameters={
        "levels": ("fixed", (252.0, 253.0)),
    },
    check_ruleset_name="p1_voltage",
)


# ---------------------------------------------------------------------------
# Current
# ---------------------------------------------------------------------------

CURRENT_ITEMS = {
    "Total": "active_current_a",
    "L1": "active_current_l1_a",
    "L2": "active_current_l2_a",
    "L3": "active_current_l3_a",
}


def discover_p1_current(section):
    for item, key in CURRENT_ITEMS.items():
        if key in section:
            yield Service(item=item)


def check_p1_current(item, section):
    key = CURRENT_ITEMS.get(item)

    if key is None or key not in section:
        yield Result(
            state=State.UNKNOWN,
            summary="Current value is not available",
        )
        return

    value = float(section[key])

    yield Result(
        state=State.OK,
        summary=f"{value:.3f} A",
    )

    yield Metric(
        name="p1_current_a",
        value=value,
        boundaries=(0, None),
    )


check_plugin_p1_current = CheckPlugin(
    name="p1_current",
    sections=["p1"],
    service_name="P1 Current %s",
    discovery_function=discover_p1_current,
    check_function=check_p1_current,
)


# ---------------------------------------------------------------------------
# Energy import / export
# ---------------------------------------------------------------------------

ENERGY_ITEMS = {
    "Import": {
        "total": "total_power_import_kwh",
        "t1": "total_power_import_t1_kwh",
        "t2": "total_power_import_t2_kwh",
    },
    "Export": {
        "total": "total_power_export_kwh",
        "t1": "total_power_export_t1_kwh",
        "t2": "total_power_export_t2_kwh",
    },
}


def discover_p1_energy(section):
    for item, keys in ENERGY_ITEMS.items():
        if any(key in section for key in keys.values()):
            yield Service(item=item)


def check_p1_energy(item, section):
    keys = ENERGY_ITEMS.get(item)

    if keys is None:
        yield Result(
            state=State.UNKNOWN,
            summary="Unknown energy type",
        )
        return

    found = False

    for metric_name, key in (
        ("p1_energy_total_kwh", keys["total"]),
        ("p1_energy_t1_kwh", keys["t1"]),
        ("p1_energy_t2_kwh", keys["t2"]),
    ):
        if key not in section:
            continue

        found = True

        yield Metric(
            name=metric_name,
            value=float(section[key]),
            boundaries=(0, None),
        )

    if not found:
        yield Result(
            state=State.UNKNOWN,
            summary=f"{item} energy values are not available",
        )
        return

    yield Result(
        state=State.OK,
        summary=f"{item} energy data available",
    )


check_plugin_p1_energy = CheckPlugin(
    name="p1_energy",
    sections=["p1"],
    service_name="P1 Energy %s",
    discovery_function=discover_p1_energy,
    check_function=check_p1_energy,
)


# ---------------------------------------------------------------------------
# Gas
# ---------------------------------------------------------------------------

def discover_p1_gas(section):
    if "total_gas_m3" in section:
        yield Service()


def check_p1_gas(section):
    value = float(section["total_gas_m3"])

    yield Result(
        state=State.OK,
        summary=f"{value:.2f} m³",
    )

    yield Metric(
        name="p1_gas_m3",
        value=value,
        boundaries=(0, None),
    )


check_plugin_p1_gas = CheckPlugin(
    name="p1_gas",
    sections=["p1"],
    service_name="P1 Gas",
    discovery_function=discover_p1_gas,
    check_function=check_p1_gas,
)


# ---------------------------------------------------------------------------
# Power quality
# ---------------------------------------------------------------------------

POWER_QUALITY_KEYS = {
    "Voltage sag L1": "voltage_sag_l1_count",
    "Voltage sag L2": "voltage_sag_l2_count",
    "Voltage sag L3": "voltage_sag_l3_count",
    "Voltage swell L1": "voltage_swell_l1_count",
    "Voltage swell L2": "voltage_swell_l2_count",
    "Voltage swell L3": "voltage_swell_l3_count",
    "Any power failure": "any_power_fail_count",
    "Long power failure": "long_power_fail_count",
}


def discover_p1_power_quality(section):
    if any(key in section for key in POWER_QUALITY_KEYS.values()):
        yield Service()


def check_p1_power_quality(section):
    found = 0

    for metric_name, key in POWER_QUALITY_KEYS.items():
        if key not in section:
            continue

        found += 1

        metric_id = (
            metric_name.lower()
            .replace(" ", "_")
        )

        yield Metric(
            name=f"p1_{metric_id}",
            value=float(section[key]),
            boundaries=(0, None),
        )

    if found:
        yield Result(
            state=State.OK,
            summary=f"{found} power-quality counters available",
        )
    else:
        yield Result(
            state=State.UNKNOWN,
            summary="No power-quality counters available",
        )


check_plugin_p1_power_quality = CheckPlugin(
    name="p1_power_quality",
    sections=["p1"],
    service_name="P1 Power Quality",
    discovery_function=discover_p1_power_quality,
    check_function=check_p1_power_quality,
)