from prometheus_client import start_http_server, Summary, Gauge
import paho.mqtt.client as mqtt
import time, json, re
from node_api import NodeAPI

client_id = "exporter"
topics = [
        ("$SYS/broker/clients/connected", 0),
        ("$CONNECTED/#", 0),
        ("+/alert", 0),
        ("homelab/ats/supply/+", 0)
        ]
nodes = {
        'rack-temp': 'Office RPi temperature probes',
        'powerpi': 'Homelab RPi power monitoring',
        'hass': 'Home Assistant',
        'exporter': 'Prometheus MQTT topic exporter'
        }

node_api = NodeAPI()

MQTT_CONNECTED = Gauge('mqtt_node_connected', 'MQTT node connection status', ['node', 'description'])
MQTT_BROKER_CLIENTS = Gauge('mqtt_broker_clients', 'MQTT broker clients connected')
MQTT_ALERT = Gauge('mqtt_alert_state', 'MQTT alert topic states', ['origin', 'alert'])
MQTT_STATUS = Gauge('mqtt_status', 'MQTT status topic', ['status'])

def get_desc(node):
    if node in nodes:
        return nodes[node]
    else:
        return node_api.get(node)

def on_log(client, userdata, level, buff):
    print(buff)

def on_connect(client, userdata, flags, rc):
    print("Connected with result code " + str(rc))

    client.subscribe(topics)
    if rc==0:
        client.connected_flag=True
        client.publish("$CONNECTED/"+client_id, 1, retain=True)
    else:
        client.bad_connection_flag=True


#def on_disconnect(client, userdata, rc):
#    client.publish(connection_topic,0,retain=True)

def on_message(client, userdata, msg):
    topic = msg.topic.split('/')

    if topic[0] == "$CONNECTED":
        MQTT_CONNECTED.labels(topic[1], get_desc(topic[1])).set(int(msg.payload))

    elif msg.topic == "$SYS/broker/clients/connected":
        MQTT_BROKER_CLIENTS.set(int(msg.payload))

    elif topic[1] == "alert":
        for alert, state in json.loads(msg.payload).items():
            MQTT_ALERT.labels(topic[0], alert).set(int(state))

    elif re.match("homelab/ats/supply/src[12]", msg.topic):
        MQTT_STATUS.labels('ats_preferred_' + topic[3]).set(json.loads(msg.payload)['preferred'])

    try:
        my_json = json.loads(msg.payload)
        print(msg.topic, json.dumps(my_json, indent=4, sort_keys=True))
    except ValueError as e:
        print(msg.topic, msg.payload)

client = mqtt.Client(client_id)

client.on_connect = on_connect
#client.on_disconnect = on_disconnect
client.on_message = on_message
#client.on_log = on_log
client.will_set("$CONNECTED/"+client_id, 0, qos=0, retain=True)
client.connect('localhost')
client.loop_start()


if __name__ == '__main__':
    # Start up the server to expose the metrics.
    start_http_server(8000)

    while True:
        time.sleep(1)