patlite-mqtt/patlite_mqtt/daemon.py

from __future__ import annotations

import argparse
import logging
import os
import signal
import sys

import paho.mqtt.client as mqtt
from usb.core import USBError

from patlite_mqtt.channels import Mode, PatliteController
from patlite_mqtt.device import PatliteDevice

LOG = logging.getLogger("patlite_mqtt")


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Patlite tower light MQTT daemon")
    parser.add_argument(
        "--broker",
        default=os.environ.get("MQTT_BROKER", "localhost"),
        help="MQTT broker host (env: MQTT_BROKER)",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=int(os.environ.get("MQTT_PORT", "1883")),
        help="MQTT broker port (env: MQTT_PORT)",
    )
    parser.add_argument(
        "--username",
        default=os.environ.get("MQTT_USERNAME"),
        help="MQTT username (env: MQTT_USERNAME)",
    )
    parser.add_argument(
        "--password",
        default=os.environ.get("MQTT_PASSWORD"),
        help="MQTT password (env: MQTT_PASSWORD)",
    )
    parser.add_argument(
        "--prefix",
        default=os.environ.get("MQTT_PREFIX", "patlite"),
        help="Topic prefix; channels are <prefix>/<color> (env: MQTT_PREFIX)",
    )
    parser.add_argument(
        "--client-id",
        default=os.environ.get("MQTT_CLIENT_ID", "patlite-mqtt"),
        help="MQTT client id (env: MQTT_CLIENT_ID)",
    )
    parser.add_argument(
        "--fast-delay",
        type=float,
        default=float(os.environ.get("FAST_FLASH_DELAY", "0.25")),
        help="Seconds per half-cycle for FAST_FLASH (env: FAST_FLASH_DELAY)",
    )
    parser.add_argument(
        "--slow-delay",
        type=float,
        default=float(os.environ.get("SLOW_FLASH_DELAY", "0.5")),
        help="Seconds per half-cycle for SLOW_FLASH (env: SLOW_FLASH_DELAY)",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Enable debug logging",
    )
    return parser.parse_args(argv)


class PatliteMqttDaemon:
    def __init__(self, args: argparse.Namespace, controller: PatliteController):
        self._args = args
        self._controller = controller
        self._client = mqtt.Client(
            mqtt.CallbackAPIVersion.VERSION2,
            client_id=args.client_id,
        )
        if args.username:
            self._client.username_pw_set(args.username, args.password or None)
        self._client.on_connect = self._on_connect
        self._client.on_message = self._on_message

    def _topic(self, color: str) -> str:
        return f"{self._args.prefix}/{color}"

    def _on_connect(self, client, userdata, flags, reason_code, properties):
        if reason_code != 0:
            LOG.error("MQTT connect failed: %s", reason_code)
            return
        LOG.info("Connected to MQTT broker")
        for color in self._controller.channels:
            topic = self._topic(color)
            client.subscribe(topic)
            LOG.info("Subscribed to %s", topic)

    def _on_message(self, client, userdata, message):
        color = message.topic.rsplit("/", 1)[-1]
        channel = self._controller.channels.get(color)
        if channel is None:
            LOG.warning("Message on unexpected topic %s", message.topic)
            return

        payload = message.payload.decode("utf-8", errors="replace").strip()
        try:
            mode = Mode.parse(payload)
        except ValueError as exc:
            LOG.warning("Ignoring %s: %s", message.topic, exc)
            return

        LOG.info("%s -> %s", message.topic, mode.value)
        channel.set_mode(mode)

    def run(self):
        self._client.connect(
            self._args.broker,
            self._args.port,
            keepalive=60,
        )
        self._client.loop_forever()

    def stop(self):
        self._client.loop_stop()
        self._client.disconnect()


def main(argv: list[str] | None = None) -> int:
    try:
        args = parse_args(argv)
    except (ValueError, TypeError) as exc:
        print(f"Invalid configuration: {exc}", file=sys.stderr)
        return 1

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )

    device = PatliteDevice()
    try:
        LOG.info("Opening Patlite USB device")
        device.connect()
    except (RuntimeError, USBError) as exc:
        LOG.error("%s", exc)
        return 1

    controller = PatliteController(
        device,
        fast_delay=args.fast_delay,
        slow_delay=args.slow_delay,
    )
    daemon = PatliteMqttDaemon(args, controller)

    def handle_signal(signum, frame):
        LOG.info("Shutting down (signal %s)", signum)
        daemon.stop()
        controller.shutdown()
        device.close()
        sys.exit(0)

    signal.signal(signal.SIGINT, handle_signal)
    signal.signal(signal.SIGTERM, handle_signal)

    try:
        LOG.info(
            "Connecting to MQTT broker %s:%s (client_id=%s)",
            args.broker,
            args.port,
            args.client_id,
        )
        daemon.run()
    except Exception as exc:
        LOG.error("MQTT connection failed: %s", exc)
        return 1
    finally:
        controller.shutdown()
        device.close()

    return 0


if __name__ == "__main__":
    raise SystemExit(main())