#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0-only
# udt-ledstrip: light a Tuya Wi-Fi LED strip in the wallpaper accent.
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
# Licensed under the GNU General Public License v2 only.
"""Run the Hyprland border's gradient along every Tuya addressable strip.

Usage: udt-ledstrip
       udt-ledstrip --selftest

The colours are the three udt-accent wrote for the window border, the accent
between its two hue neighbours, played as the strip's "flow" scene: a gradient
moving along the strip. udt-accent starts this detached on every run, like
udt-keyboard.

The device id and local key come from ~/.config/udt/tuya/devices.json, written
once by `python -m tinytuya wizard` against a Tuya IoT cloud project; after
that everything is local. Re-pairing a light in the app changes its key, so
the wizard has to run again. The address is not stored: the light broadcasts
it on UDP 6667 every few seconds, so a DHCP change costs nothing.

Speaks Tuya protocol 3.5 itself (AES-GCM, session key negotiated per
connection) with the `cryptography` module, rather than depending on tinytuya.
Without the device list, or with the light unplugged, this does nothing.
"""
import base64
import colorsys
import hashlib
import hmac
import json
import os
import re
import socket
import struct
import sys
import time
from pathlib import Path

DEVICES = Path.home() / ".config/udt/tuya/devices.json"
BORDER = Path.home() / ".cache/udt/udt-border.lua"

# Every Tuya device encrypts its discovery broadcast with this published key.
UDP_KEY = hashlib.md5(b"yGAdlopoPVldABfn").digest()
DISCOVERY_TIMEOUT = 20

# Calibration, all 0..100. LEDs render the schemes' pastel accents as near
# white, so saturation is floored: 100 plays every hue pure, lower lets a
# pale accent stay paler than a vivid one. Brightness is fixed rather than
# read back from the strip, so it overrides whatever the app last set.
SATURATION_FLOOR = 100
BRIGHTNESS = 100
SPEED = 50

SESS_START, SESS_FINISH, CONTROL = 3, 5, 13

# DP 51, dreamlight_scene_mode, as the app writes it (read off the strip's
# status pushes; Tuya does not document it): version 1, scene id (0 for a
# custom one), mode, speed twice, flags, two zero bytes, brightness, then
# per colour a two byte hue and a saturation byte. Mode 0x0a with flags
# 0xe0 is the app's moving gradient.
FLOW, FLOW_FLAGS = 0x0A, 0xE0


def scene(colours):
    """DP 51 bytes playing #rrggbb colours as a flow, base64 for the wire."""
    body = bytearray([1, 0, FLOW, SPEED, SPEED, FLOW_FLAGS, 0, 0, BRIGHTNESS])
    for c in colours:
        h, s, _ = colorsys.rgb_to_hsv(*(int(c[i:i + 2], 16) / 255 for i in (1, 3, 5)))
        body += (round(h * 360) % 360).to_bytes(2, "big")
        body.append(max(round(s * 100), SATURATION_FLOOR))
    return base64.b64encode(body).decode()


def pack(key, seq, cmd, payload):
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    iv = os.urandom(12)
    header = struct.pack(">IHIII", 0x6699, 0, seq, cmd, 12 + len(payload) + 16)
    return header + iv + AESGCM(key).encrypt(iv, payload, header[4:]) + b"\0\0\x99\x66"


def unpack(key, msg):
    """(cmd, plaintext) of one 0x6699 frame."""
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    _, _, _, cmd, length = struct.unpack(">IHIII", msg[:18])
    body = msg[18:18 + length]
    return cmd, AESGCM(key).decrypt(body[:12], body[12:], msg[4:18])


def recv_frame(sock):
    def exact(n):
        buf = b""
        while len(buf) < n:
            chunk = sock.recv(n - len(buf))
            if not chunk:
                raise ConnectionError("light closed the connection")
            buf += chunk
        return buf
    header = exact(18)
    return header + exact(struct.unpack(">I", header[14:18])[0] + 4)


def discover(ids):
    """{id: ip} for the wanted ids heard broadcasting within the timeout."""
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
    sock.bind(("", 6667))
    found, deadline = {}, time.monotonic() + DISCOVERY_TIMEOUT
    while set(found) != ids and (left := deadline - time.monotonic()) > 0:
        sock.settimeout(left)
        try:
            data, _ = sock.recvfrom(4096)
            # Device-sent payloads start with a four byte return code.
            info = json.loads(unpack(UDP_KEY, data)[1][4:])
        except Exception:  # timeout, or not a 3.5 frame (InvalidTag)
            continue
        if info.get("gwId") in ids:
            found[info["gwId"]] = info["ip"]
    return found


def set_dps(ip, key, dps):
    real_key = key.encode()
    with socket.create_connection((ip, 6668), timeout=5) as sock:
        local_nonce = os.urandom(16)
        sock.sendall(pack(real_key, 1, SESS_START, local_nonce))
        reply = unpack(real_key, recv_frame(sock))[1][-48:]
        remote_nonce, mac = reply[:16], reply[16:]
        if not hmac.compare_digest(mac, hmac.digest(real_key, local_nonce, "sha256")):
            raise ValueError("handshake failed: local key is wrong or stale")
        sock.sendall(pack(real_key, 2, SESS_FINISH,
                          hmac.digest(real_key, remote_nonce, "sha256")))

        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        mixed = bytes(a ^ b for a, b in zip(local_nonce, remote_nonce))
        session = AESGCM(real_key).encrypt(local_nonce[:12], mixed, None)[:16]

        body = json.dumps({"protocol": 5, "t": int(time.time()),
                           "data": {"dps": dps}}, separators=(",", ":"))
        sock.sendall(pack(session, 3, CONTROL, b"3.5" + b"\0" * 12 + body.encode()))
        unpack(session, recv_frame(sock))  # the ack; raises if it is garbage


def selftest():
    raw = base64.b64decode(scene(["#ff0000", "#0000ff"]))
    assert raw[:9] == bytes([1, 0, FLOW, SPEED, SPEED, FLOW_FLAGS, 0, 0, BRIGHTNESS])
    assert raw[9:] == bytes([0, 0, 100, 0, 240, 100])
    key = b"0123456789abcdef"
    assert unpack(key, pack(key, 7, CONTROL, b"payload")) == (CONTROL, b"payload")
    print("ok")


def main():
    if sys.argv[1:] == ["--selftest"]:
        return selftest()
    if sys.argv[1:]:
        sys.exit(__doc__.split("\n\n")[1])
    try:
        devices = json.loads(DEVICES.read_text())
    except (OSError, ValueError):
        return
    # Lights only: anything else the Tuya account links has no colour DP.
    # The cloud spec omits DP 51, so this matches on the DP 24 it does list.
    lights = {d["id"]: d["key"] for d in devices
              if d.get("mapping", {}).get("24", {}).get("code") == "colour_data"}
    border = BORDER.read_text() if BORDER.exists() else ""
    colours = ["#" + h for h in re.findall(r"rgb\(([0-9a-fA-F]{6})\)", border)]
    if not colours or not lights:
        return

    dps = {"21": "scene", "51": scene(colours)}
    for dev_id, ip in discover(set(lights)).items():
        try:
            set_dps(ip, lights[dev_id], dps)
        except Exception as e:
            print(f"udt-ledstrip: {dev_id} at {ip}: {e}", file=sys.stderr)


if __name__ == "__main__":
    main()