#!/bin/env python3
import asyncio
from pysnmp.hlapi.asyncio.slim import Slim
from pysnmp.smi.rfc1902 import ObjectIdentity, ObjectType
import logging
import random
from paho.mqtt import client as mqtt_client
import json
from time import sleep

logging.basicConfig(
        format='(%(levelname)s) %(message)s',
        level=logging.DEBUG
        )


def connect_mqtt(mqtt_config):
    def on_connect(client, userdata, flags, rc):
        if rc == 0:
            print("Connected to MQTT Broker!")
        else:
            print("Failed to connect, return code {rc}")
    client = mqtt_client.Client()
    client.username_pw_set(mqtt_config["user"], mqtt_config["password"])
    client.on_connect = on_connect
    client.connect(mqtt_config['broker'], mqtt_config['port'])
    return client


def publish(topic, client, data, retain, qos):
    msg = json.dumps(data)
    result = client.publish(topic=topic, payload=msg, qos=qos, retain=bool(retain))
    status = result[0]
    if status == 0:
        logging.debug(f"Send `{msg}` to topic `{topic}`")
    else:
        logging.error(f"Failed to send message to topic {topic}")


async def get_snmp(req):
    data = {}
    for oid in req["oids"]:
        with Slim(1) as slim:
            errorIndication, errorStatus, errorIndex, varBinds = await slim.get(
                req["snmp_community"],
                req["ip"],
                161,
                ObjectType(ObjectIdentity(oid["oid"])),
            )

            if errorIndication:
                logging.error(errorIndication)
            elif errorStatus:
                logging.error(
                    "{} at {}".format(
                        errorStatus.prettyPrint(),
                        errorIndex and varBinds[int(errorIndex) - 1][0] or "?",
                    )
                )
            else:
                for varBind in varBinds:
                    logging.debug(f"{req['device_name']} {oid['name']} => {oid['type'](varBind[1])}")
                    if oid['type'] == bool:
                        if bool(varBind[1]):
                            data.update({oid["name"]: "on"})
                        else:
                            data.update({oid["name"]: "off"})
                    else:
                        data.update({oid["name"]: oid["type"](varBind[1])})
    logging.debug(f"JSON : {json.dumps(data)}")
    return data


def ha_create_config(req):
    ha_config = {}
    device = {
            "ids": f"{req['device_name']}_{req['ip']}".replace(".", "_"),
            "name": req['device_name'],
            }
    origin = {
            "name": "snmp2mqtt"
            }
    ha_config.update({"dev": device, "o": origin})
    ha_config.update({"state_topic": f"SNMP/{req['device_name']}/state"})
    ha_config.update({"qos": 1})
    cmps = {}
    for oid in req['oids']:
        cmps.update(
                {
                    f"{req['device_name']}_{req['ip']}_{oid['name']}".replace(".", "_"):
                    {
                        "p": oid['HA_platform'],
                        "device_class": oid['HA_device_class'],
                        "value_template": f"{{{{ value_json.{oid['name']}}}}}",
                        "unique_id": f"{req['device_name']}_{req['ip']}_{oid['name']}".replace(".", "_"),
                        "name": oid['name']
                        }
                    })
        if "HA_unit" in oid.keys():
            cmps.update(
                 {f"{req['device_name']}_{req['ip']}_{oid['name']}".replace(".", "_"):
                    {"unit_of_measurement": oid['HA_unit']}})
    ha_config.update({"cmps": cmps})
    logging.debug(f"config : {json.dumps(ha_config)}")
    return ha_config


def send_to_mqtt():
    config = ha_create_config(req)
    client = connect_mqtt(mqtt_config)
    client.loop_start()
    config_topic = f"homeassistant/device/{config['dev']['ids']}/config"
    state_topic = config['state_topic']
    while True:
        try:
            publish(config_topic, client, config, True, 2)
            logging.info(f"{config_topic} -> {config}")
        except Exception as e:
            logging.error(e)
            pass
        try:
            state = asyncio.run(get_snmp(req))
            publish(state_topic, client, state, False, 0)
            logging.info(f"{state_topic} -> {state}")
        except Exception as e:
            logging.error(e)
            pass
        sleep(2)


req = {
        "device_name": "mikrotik_hex",
        "ip": "192.168.10.2",
        "snmp_community": "public",
        "oids": [
            {"name": "stln_vpn_in",
             "oid": ".1.3.6.1.2.1.2.2.1.10.12",
             "type": int,
             "HA_device_class": "data_size",
             "HA_platform": "sensor",
             "HA_unit": "bit"
             },
            {"name": "stlon_vpn_out",
             "oid": ".1.3.6.1.2.1.2.2.1.16.12",
             "type": int,
             "HA_device_class": "data_size",
             "HA_platform": "sensor",
             "HA_unit": "bit"
             },
            {"name": "stln_vpn_status",
             "oid": ".1.3.6.1.2.1.2.2.1.8.12",
             "type": bool,
             "HA_device_class": "connectivity",
             "HA_platform": "binary_sensor",
             }
            ]
        }

mqtt_config = {
        "broker": "192.168.10.202",
        "port": 1883,
        "client_id": f"snmp-mqtt-{random.randint(0, 1000)}",
        "user": "snmp2mqtt",
        "password": "snmp_2_MQTT"
        }


send_to_mqtt()