import asyncio
import json
import logging

from aiomqtt import Client
from prometheus_client import Counter, Gauge
from rdferry import StarletteServer

logging.basicConfig(level=logging.DEBUG)
log = logging.getLogger()

powermeter_reports = Counter('powermeter_reports', 'mqtt messages from powermeter', labelnames=['sensor'])
metrics = {}
for metric_name in [
    ('powermeter_w'),
    ('powermeter_apparent_power_va'),
    ('powermeter_reactive_power_var'),
    ('powermeter_factor'),
    ('powermeter_voltage_v'),
    ('powermeter_current_a'),
    ('powermeter_temp_c'),
]:
    metrics[metric_name] = Gauge(metric_name, documentation='', labelnames=['sensor'])


def sonoff(sensor_label, msg):
    energy = msg['ENERGY']
    metrics['powermeter_w'].labels(sensor=sensor_label).set(energy['Power'])
    metrics['powermeter_apparent_power_va'].labels(sensor=sensor_label).set(energy['ApparentPower'])
    metrics['powermeter_reactive_power_var'].labels(sensor=sensor_label).set(energy['ReactivePower'])
    metrics['powermeter_factor'].labels(sensor=sensor_label).set(energy['Factor'])
    metrics['powermeter_voltage_v'].labels(sensor=sensor_label).set(energy['Voltage'])
    metrics['powermeter_current_a'].labels(sensor=sensor_label).set(energy['Current'])


def shelly(sensor_label, msg):
    metrics['powermeter_w'].labels(sensor=sensor_label).set(msg['apower'])
    metrics['powermeter_voltage_v'].labels(sensor=sensor_label).set(msg['voltage'])
    metrics['powermeter_current_a'].labels(sensor=sensor_label).set(msg['current'])
    metrics['powermeter_apparent_power_va'].labels(sensor=sensor_label).set(msg['apower'])
    metrics['powermeter_temp_c'].labels(sensor=sensor_label).set(msg['temperature']['tC'])


async def subscribe_to_sensor(topic, sensor_label, onMsg):
    async with Client("mqtt2") as client:
        await client.subscribe(topic)
        async for message in client.messages:
            try:
                msg = json.loads(message.payload)
                onMsg(sensor_label, msg)
                powermeter_reports.labels(sensor=sensor_label).inc()
            except Exception:
                log.error(f'{topic=} {sensor_label=}', exc_info=True)


async def request_shelly_statuses(topic, period=10):
    async with Client("mqtt2") as client:
        while True:
            # todo: this adds lag- better to publish upon a /metrics request
            # then block until all the statuses come in.
            await client.publish(topic, 'status_update')
            await asyncio.sleep(period)


server = StarletteServer()
server.serve(startup_tasks=[
    subscribe_to_sensor(    topic="tele/st-wall-power/SENSOR",       sensor_label='st_wall',    onMsg=sonoff),
    subscribe_to_sensor(    topic="tele/ki-fridge/SENSOR",           sensor_label='ki_fridge',  onMsg=sonoff),
    subscribe_to_sensor(    topic="tt-console-power/SENSOR",         sensor_label='tt_console', onMsg=sonoff),
    request_shelly_statuses(topic="ga-washer-power/command/switch:0"),
    subscribe_to_sensor(    topic="ga-washer-power/status/switch:0", sensor_label='ga_washer',  onMsg=shelly),
    request_shelly_statuses(topic="do-r-power/command/switch:0"),
    subscribe_to_sensor(    topic="do-r-power/status/switch:0",      sensor_label='do_r',       onMsg=shelly),
    request_shelly_statuses(topic="ws-atx-power/command/switch:0"),
    subscribe_to_sensor(    topic="ws-atx-power/status/switch:0",    sensor_label='ws_atx',     onMsg=shelly),
])  # yapf: disable