import json

import time

import paho.mqtt.client as mqtt

from prometheus_client import start_http_server, Gauge

CONFIG = json.load(open('/etc/mqtt2prometheus.json', mode='r', encoding='utf-8'))

people_now_present = Gauge('people_now_present', 'People currently present in the space')

space_state = Gauge('space_state', 'The current space state')

space_lastchange = Gauge('space_lastchange', 'The lastchange timestamp of the current space state')

member_count = Gauge('member_count', 'The current count of members', labelnames=['name'])

availability = Gauge('availability', 'The availability status of different MQTT clients', labelnames=['clientname'])

def on_connect(client, userdata, flags, rc):

    client.publish('availability/mqtt2prometheus', 'online', 2, True)

    client.message_callback_add('sensors/people_now_present', handle_people_now_present)

    client.message_callback_add('space/status', handle_space_status)

    client.message_callback_add('space/member_count', handle_member_count)

    client.message_callback_add('availability/#', handle_availability)

    client.subscribe('sensors/people_now_present')

    client.subscribe('space/status')

    client.subscribe('space/member_count')

    client.subscribe('availability/+')

    client.subscribe('availability/shutdown/+')

def handle_people_now_present(client, userdata, msg):

    people_now_present.set(json.loads(msg.payload))

def handle_space_status(client, userdata, msg):

    state = json.loads(msg.payload)

    space_state.set(int(state['open']))

    space_lastchange.set(state['lastchange'])

def handle_member_count(client, userdata, msg):

    count = json.loads(msg.payload)

    member_count.labels(count['name']).set(count['value'])

def handle_availability(client, userdata, msg):

    t = msg.topic.split('/')

    if len(t) == 2:

        c = t[-1]

    else:

        c = '_'.join(t[1:])

    v = 1 if msg.payload == b'online' else 0

    availability.labels(c).set(v)

def main():

    client = mqtt.Client('mqtt2prometheus')

    client.on_connect = on_connect

    client.username_pw_set(CONFIG['mqtt']['username'], CONFIG['mqtt']['password'])

    client.will_set('availability/mqtt2prometheus', 'offline', 2, True)

    client.connect(CONFIG['mqtt']['host'])

    client.loop_start()

    start_http_server(CONFIG['prometheus']['port'])

    while True:

        time.sleep(1)

if __name__ == '__main__':

    main()