#!/usr/bin/env python3
"""MQTT diagnostic helpers — invoked by justfile mqtt-* recipes."""
import argparse
import sys
import time

import paho.mqtt.client as mqtt


def pub(host, port, topic, payload):
    c = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
    c.connect(host, port, 60)
    c.loop_start()
    info = c.publish(topic, payload)
    info.wait_for_publish(timeout=5)
    c.loop_stop()
    c.disconnect()


def sub(host, port, topic, seconds):
    got = []

    def on_msg(client, userdata, msg):
        got.append((msg.topic, msg.payload.decode()))

    c = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
    c.on_message = on_msg
    c.connect(host, port, 60)
    c.loop_start()
    c.subscribe(topic)
    end = time.monotonic() + seconds
    while time.monotonic() < end:
        time.sleep(0.1)
    c.loop_stop()
    c.disconnect()
    for t, p in got:
        print(t, p)


def watch(host, port, topic):
    def on_msg(client, userdata, msg):
        print(msg.topic, msg.payload.decode())

    c = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
    c.on_message = on_msg
    c.connect(host, port, 60)
    c.loop_start()
    c.subscribe(topic)
    try:
        while True:
            time.sleep(0.5)
    except KeyboardInterrupt:
        pass
    c.loop_stop()
    c.disconnect()


def main():
    ap = argparse.ArgumentParser()
    sub_p = ap.add_subparsers(dest="cmd", required=True)

    ppub = sub_p.add_parser("pub", help="Publish a single message")
    ppub.add_argument("--host", default="192.168.81.147")
    ppub.add_argument("--port", type=int, default=1883)
    ppub.add_argument("--topic", required=True)
    ppub.add_argument("--payload", required=True)

    psub = sub_p.add_parser("sub", help="Subscribe for N seconds and print messages")
    psub.add_argument("--host", default="192.168.81.147")
    psub.add_argument("--port", type=int, default=1883)
    psub.add_argument("--topic", required=True)
    psub.add_argument("--seconds", type=float, default=5.0)

    pwatch = sub_p.add_parser("watch", help="Subscribe forever and print messages")
    pwatch.add_argument("--host", default="192.168.81.147")
    pwatch.add_argument("--port", type=int, default=1883)
    pwatch.add_argument("--topic", default="zigbee2mqtt/#")

    args = ap.parse_args()
    if args.cmd == "pub":
        pub(args.host, args.port, args.topic, args.payload)
    elif args.cmd == "sub":
        sub(args.host, args.port, args.topic, args.seconds)
    elif args.cmd == "watch":
        watch(args.host, args.port, args.topic)


if __name__ == "__main__":
    main()