#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0-only
# udt-bulb: light IKEA TRADFRI colour bulbs in the wallpaper accent.
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
# Licensed under the GNU General Public License v2 only.
"""Fade every colour bulb on a TRADFRI gateway to the wallpaper accent.

Usage: udt-bulb                  show the accent, leaving the bulbs on or off
       udt-bulb on|off           switch them, the accent with them when on
       udt-bulb --pair HOST CODE pair with the gateway, list its bulbs
       udt-bulb --selftest

The verbs are udt-ledstrip's, and its callers call this beside it, except on
suspend: the strip breathes to say the desktop is asleep, and the bulb is
simply switched off. A bulb shows one colour, so it gets the accent alone: the
middle of the three udt-accent wrote for the window border. A bare run touches only bulbs that are
already on, because the gateway switches a bulb on for any colour change.

The bulbs are Zigbee and have no address; the gateway takes commands for them
over CoAP with DTLS-PSK on UDP 5684, which Python cannot speak, so this calls
libcoap's coap-client. `--pair` trades the security code on the gateway's
label for an identity and key of its own, and records the gateway and its
colour bulbs in ~/.config/udt/tradfri.json; pair again after adding a bulb.
Without that file, or without coap-client, this does nothing.
"""
import colorsys
import json
import re
import secrets
import shutil
import subprocess
import sys
from pathlib import Path

CONFIG = Path.home() / ".config/udt/tradfri.json"
BORDER = Path.home() / ".cache/udt/udt-border.lua"

# Calibration. Hue and saturation run 0..65279 on the gateway, brightness
# 0..254, transition in tenths of a second. Bulbs wash the schemes' pastel
# accents out to white, as the LED strip did, so saturation is floored.
SATURATION_FLOOR = 1.0
BRIGHTNESS = 254
FADE = 20
TIMEOUT = 5

# IPSO object ids the gateway uses: 15001 devices, 3311 light control,
# 5850 on/off, 5851 brightness, 5707 hue, 5708 saturation, 5712 transition.
DEVICES, LIGHT = "15001", "3311"
# The gateway drops part of a request that sets brightness and colour at
# once (both stalled mid-fade), and a colour sent while a brightness fade
# runs, so colour goes first and brightness after, in requests of their own.
HUE_MAX = 65279


def coap(method, url, ident, key, payload=None):
    """The response body of one coaps request, or None on any failure."""
    cmd = ["coap-client", "-m", method, "-u", ident, "-k", key, "-B", str(TIMEOUT)]
    if payload is not None:
        cmd += ["-e", json.dumps(payload, separators=(",", ":"))]
    try:
        out = subprocess.run(cmd + [url], capture_output=True, text=True,
                             timeout=TIMEOUT + 2, check=True).stdout
    except (subprocess.SubprocessError, OSError):
        return None
    return out.strip()


def light_states(colour):
    """3311 states for #rrggbb, sent in order."""
    h, s, _ = colorsys.rgb_to_hsv(*(int(colour[i:i + 2], 16) / 255 for i in (1, 3, 5)))
    s = max(s, SATURATION_FLOOR)
    return [{"5707": round(h * HUE_MAX), "5708": round(s * HUE_MAX), "5712": FADE},
            {"5850": 1, "5851": BRIGHTNESS, "5712": FADE}]


def pair(host, code):
    ident = "udt-" + secrets.token_hex(4)
    base = f"coaps://{host}:5684"
    reply = coap("post", f"{base}/15011/9063", "Client_identity", code, {"9090": ident})
    try:
        key = json.loads(reply)["9091"]
    except (TypeError, ValueError, KeyError):
        sys.exit(f"udt-bulb: pairing with {host} failed: {reply or 'no reply'}")
    bulbs = []
    for dev in json.loads(coap("get", f"{base}/{DEVICES}", ident, key) or "[]"):
        info = json.loads(coap("get", f"{base}/{DEVICES}/{dev}", ident, key) or "{}")
        # Colour bulbs report a hue; white-spectrum ones only a colour temperature.
        if "5707" in (info.get(LIGHT) or [{}])[0]:
            bulbs.append(dev)
            print(f"{dev}: {info.get('9001', '?')}")
    CONFIG.parent.mkdir(parents=True, exist_ok=True)
    CONFIG.write_text(json.dumps({"host": host, "identity": ident, "key": key,
                                  "bulbs": bulbs}))
    CONFIG.chmod(0o600)
    print(f"{len(bulbs)} colour bulb(s) recorded in {CONFIG}")


def selftest():
    assert light_states("#ff0000") == [
        {"5707": 0, "5708": HUE_MAX, "5712": FADE},
        {"5850": 1, "5851": BRIGHTNESS, "5712": FADE}]
    # A pastel keeps its hue and loses its paleness to the floor.
    colour = light_states("#8080ff")[0]
    assert (colour["5707"], colour["5708"]) == (round(HUE_MAX * 2 / 3), HUE_MAX)
    print("ok")


def main():
    args = sys.argv[1:]
    if args == ["--selftest"]:
        return selftest()
    if len(args) == 3 and args[0] == "--pair":
        return pair(*args[1:])
    if args not in ([], ["on"], ["off"]):
        sys.exit(__doc__.split("\n\n")[1])
    try:
        cfg = json.loads(CONFIG.read_text())
    except (OSError, ValueError):
        return
    if not cfg.get("bulbs") or not shutil.which("coap-client"):
        return
    border = BORDER.read_text() if BORDER.exists() else ""
    colours = ["#" + h for h in re.findall(r"rgb\(([0-9a-fA-F]{6})\)", border)]
    accent = colours[len(colours) // 2] if colours else None

    if args == ["off"]:
        states = [{"5850": 0}]
    elif accent:
        states = light_states(accent)
    elif args == ["on"]:
        states = [{"5850": 1}]
    else:
        return
    auth = cfg["identity"], cfg["key"]
    for dev in cfg["bulbs"]:
        url = f"coaps://{cfg['host']}:5684/{DEVICES}/{dev}"
        if not args:
            # A bare run follows the wallpaper and leaves an off bulb off.
            try:
                if json.loads(coap("get", url, *auth))[LIGHT][0]["5850"] != 1:
                    continue
            except (TypeError, ValueError, KeyError, IndexError):
                print(f"udt-bulb: bulb {dev} did not answer", file=sys.stderr)
                continue
        for state in states:
            if coap("put", url, *auth, {LIGHT: [state]}) is None:
                print(f"udt-bulb: bulb {dev} did not answer", file=sys.stderr)
                break


if __name__ == "__main__":
    main()