#!/bin/python3
#
import argparse
import json
import nagios
import os
import urllib.request

def get_shelly_data(hostname):
    with urllib.request.urlopen(f"http://{hostname}/status") as response:
        return json.load(response)

def check_wifi(hostname, warn_lvl, crit_lvl):
    data = get_shelly_data(hostname)
    #
    connected = data['wifi_sta']['connected']
    quality = data['wifi_sta']['rssi']
    if not connected or quality <= crit_lvl:
        status = nagios.Nagios.ERROR
    elif quality <= warn_lvl:
        status = nagios.Nagios.WARNING
    else:
        status = nagios.Nagios.OK
    return (status, f"shelly: connected={connected}; quality={quality} shall be > warn={warn_lvl} / crit={crit_lvl}")

def check_mqtt(hostname, warn_lvl, crit_lvl):
    data = get_shelly_data(hostname)
    #
    connected = data['mqtt']['connected']
    if not connected:
        status = nagios.Nagios.ERROR
    else:
        status = nagios.Nagios.OK
    return (status, f"shelly: connected={connected}")

def check_memory(hostname, warn_lvl, crit_lvl):
    data = get_shelly_data(hostname)
    #
    ram_total = data['ram_total']
    ram_free = data['ram_free']
    ram_left = ram_free / ram_total * 100
    if ram_left <= crit_lvl:
        status = nagios.Nagios.ERROR
    elif ram_left <= warn_lvl:
        status = nagios.Nagios.WARNING
    else:
        status = nagios.Nagios.OK
    return (status, f"shelly: ram_left={ram_left:.1f}percent shall be > warn={warn_lvl} / crit={crit_lvl}")

def check_filesystem(hostname, warn_lvl, crit_lvl):
    data = get_shelly_data(hostname)
    #
    fs_size = data['fs_size']
    fs_free = data['fs_free']
    fs_left = fs_free / fs_size * 100
    if fs_left <= crit_lvl:
        status = nagios.Nagios.ERROR
    elif fs_left <= warn_lvl:
        status = nagios.Nagios.WARNING
    else:
        status = nagios.Nagios.OK
    return (status, f"shelly: fs_left={fs_left:.1f}percent shall be > warn={warn_lvl} / crit={crit_lvl}")

def check_temperature(hostname, warn_lvl, crit_lvl):
    data = get_shelly_data(hostname)
    #
    temperature = data.get('tmp', {}).get('tC')
    if temperature is None:
        if args.hostname.startswith('shelly1-'):
            return (nagios.Nagios.OK, "shelly: Shelly1 has no temperature information")
        status = nagios.Nagios.UNKNOWN
    elif temperature >= crit_lvl:
        status = nagios.Nagios.ERROR
    elif temperature >= warn_lvl:
        status = nagios.Nagios.WARNING
    else:
        status = nagios.Nagios.OK
    return (status, f"shelly: temperature:={temperature} C shall be < warn={warn_lvl} / crit={crit_lvl}")


if __name__ == "__main__":
    # Parse arguments
    parser = argparse.ArgumentParser(
        prog=os.path.basename(__file__),
        description='Check shelly for nagios monitorin',
    )
    parser.add_argument('-H', '--hostname', required=True)
    parser.add_argument('-w', '--warn_lvl', type=float, required=True)
    parser.add_argument('-c', '--critical_lvl', type=float, required=True)
    args = parser.parse_args()

    # Init nagios instance
    n = nagios.Nagios()

    # Identify check
    this_check = parser.prog[len("check_shelly_"):]
    check_function = {
        'wifi': check_wifi,
        'mqtt': check_mqtt,
        'memory': check_memory,
        'filesystem': check_filesystem,
        'temperature': check_temperature
    }.get(this_check)

    # process check and return feedback
    if check_function is None:
        status = n.UNKNOWN
        message = f"Unknown check - {this_check}"
    else:
        status, message = check_function(args.hostname, args.warn_lvl, args.critical_lvl)
    n.exit(status, message)