#!/usr/bin/python3
"""qmk-lights.py — QMK/VIA backlight controller for Keychron K3 Max RGB."""

import argparse
import ctypes
import ctypes.util
import dbus
import dbus.mainloop.glib
import gi.repository.GLib  # type: ignore[import-untyped]
import pprint
import sys
import time

# Must preload libhidapi-hidraw before importing hid.
# The Debian default (libhidapi-libusb) does not populate usage_page/usage,
# which are required to identify the VIA Raw HID interface.
_hidraw = ctypes.util.find_library("hidapi-hidraw")
if _hidraw:
    ctypes.CDLL(_hidraw, mode=ctypes.RTLD_GLOBAL)

import hid  # noqa: E402  # type: ignore[import-untyped]

# ── USB / HID ──────────────────────────────────────────────────────────────────
VID = 0x3434
PID = 0x0A30
RAW_HID_USAGE_PAGE = 0xFF60  # VIA Raw HID interface
RAW_HID_USAGE = 0x61
REPORT_SIZE = 32  # payload bytes (firmware sees these)

# ── VIA protocol command IDs (quantum/via.h) ───────────────────────────────────
VIA_CUSTOM_SET_VALUE = 0x07  # id_custom_set_value
VIA_CUSTOM_GET_VALUE = 0x08  # id_custom_get_value
VIA_CUSTOM_SAVE = 0x09  # id_custom_save

# ── VIA channel IDs (via.h: via_channel_id enum) ──────────────────────────────
RGB_MATRIX_CHANNEL = 0x03  # id_qmk_rgb_matrix_channel

# ── QMK RGB Matrix value IDs (via.h: via_qmk_rgb_matrix_value enum) ───────────
QMK_BRIGHTNESS = 0x01  # id_qmk_rgb_matrix_brightness
QMK_EFFECT = 0x02  # id_qmk_rgb_matrix_effect
QMK_SPEED = 0x03  # id_qmk_rgb_matrix_effect_speed
QMK_COLOR = 0x04  # id_qmk_rgb_matrix_color (hue byte, saturation byte)

# ── Preset defaults ────────────────────────────────────────────────────────────
# --on  : solid white
ON_EFFECT = 1  # RGB_MATRIX_SOLID_COLOR
ON_BRIGHTNESS = 128
ON_HUE = 0  # irrelevant when saturation = 0
ON_SATURATION = 0  # 0 = white; 255 = fully saturated colour
ON_SPEED = 128  # mid speed (unused for solid, but set anyway)

# --per-key : per_key_rgb — effect 23 (confirmed via VIA web app effect list).
PERKEY_EFFECT = 23  # per_key_rgb
PERKEY_BRIGHTNESS = 127
PERKEY_HUE = 0
PERKEY_SATURATION = 0
PERKEY_SPEED = 128

# --program : reactive_multiwide — effect 19 (confirmed via VIA web app effect list).
PROGRAMMED_EFFECT = 19  # reactive_multiwide
PROGRAMMED_BRIGHTNESS = 255
PROGRAMMED_HUE = 0
PROGRAMMED_SATURATION = 0  # full saturation makes reactive effects visible
PROGRAMMED_SPEED = 128

# ── Effect name table — confirmed from VIA web app effect list ─────────────────
EFFECT_NAMES = {
    0: "none",
    1: "solid_color",
    2: "breathing",
    3: "band_spiral_val",
    4: "cycle_all",
    5: "cycle_left_right",
    6: "cycle_up_down",
    7: "rainbow_moving_chevron",
    8: "cycle_out_in",
    9: "cycle_out_in_dual",
    10: "cycle_pinwheel",
    11: "cycle_spiral",
    12: "dual_beacon",
    13: "rainbow_beacon",
    14: "jellybean_raindrops",
    15: "pixel_rain",
    16: "typing_heatmap",
    17: "digital_rain",
    18: "reactive_simple",
    19: "reactive_multiwide",
    20: "reactive_multinexus",
    21: "splash",
    22: "solid_splash",
    23: "per_key_rgb",
    24: "mix_pgb",
}

# ── Screensaver DBus interfaces ────────────────────────────────────────────────
SCREENSAVER_INTERFACES = [
    "org.gnome.ScreenSaver",
    "org.cinnamon.ScreenSaver",
    "org.kde.screensaver",
    "org.freedesktop.ScreenSaver",
]

HELP_EPILOG = """\
━━━ Effect IDs ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

All IDs confirmed from the VIA web app effect list for this firmware.

 ID   Name
───  ──────────────────────────────────────────────────────
  0   none                    All LEDs off
  1   solid_color             Solid colour  (see --hue, --saturation)
  2   breathing               Whole-board pulsing brightness
  3   band_spiral_val         Spiral brightness band
  4   cycle_all               All keys cycle through hues together
  5   cycle_left_right        Hue sweeps left → right
  6   cycle_up_down           Hue sweeps top → bottom
  7   rainbow_moving_chevron  Chevron-shaped rainbow wave
  8   cycle_out_in            Hue cycle from edges inward
  9   cycle_out_in_dual       Dual cycle from edges inward
 10   cycle_pinwheel          Pinwheel of cycling hues
 11   cycle_spiral            Spiral of cycling hues
 12   dual_beacon             Two-colour spinning beacon
 13   rainbow_beacon          Rainbow spinning beacon
 14   jellybean_raindrops     Coloured jellybean drops
 15   pixel_rain              Pixel colour rain
 16   typing_heatmap          Keys glow where you type most
 17   digital_rain            Matrix-style digital rain
 18   reactive_simple         Solid + flash on keypress
 19   reactive_multiwide      Wide ripple on keypress
 20   reactive_multinexus     Nexus ripple on keypress
 21   splash                  Colour splash on keypress
 22   solid_splash            Solid + colour splash on keypress
 23   per_key_rgb             Per-key RGB (stored layout)
 24   mix_pgb                 Mixed PGB pattern

━━━ Colour (HSV) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

--hue 0–255 maps to the colour wheel:
  0=red  21=orange  43=yellow  64=chartreuse  85=green
  106=spring_green  128=cyan  149=azure  170=blue  192=violet
  213=magenta  234=rose  255=red (wraps)

--saturation:  0 = white (no colour),  255 = fully saturated colour.

Note: solid white is hue=<any>, saturation=0.
      To get solid blue: --effect 1 --hue 170 --saturation 255

━━━ Examples ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

  qmk-lights.py --on                    # solid white at full brightness
  qmk-lights.py --on --brightness 128   # solid white at half brightness
  qmk-lights.py --on --hue 170 --saturation 255   # solid blue
  qmk-lights.py --effect 13             # cycle_left_right (if compiled in)
  qmk-lights.py --brightness 200 --speed 64  # dim and slow, keep current effect
  qmk-lights.py --off                   # turn off, saved to EEPROM
  qmk-lights.py --query                 # read current settings from keyboard
  qmk-lights.py --screensaver           # auto off/on with screen lock
"""

debug = 0


# ── Device access ──────────────────────────────────────────────────────────────


def find_device():
    """Enumerate HID devices and return an opened handle to the VIA Raw HID interface."""
    for info in hid.enumerate(VID, 0):
        if info["product_id"] != PID:
            continue
        if info.get("usage_page") != RAW_HID_USAGE_PAGE:
            continue
        if info.get("usage") != RAW_HID_USAGE:
            continue
        if debug > 0:
            print(
                "Found: {} (usage_page=0x{:04x} usage=0x{:02x})".format(
                    info.get("path"), info.get("usage_page"), info.get("usage")
                )
            )
        kbd = hid.device()
        kbd.open_path(info["path"])
        if debug > 0:
            print("Manufacturer: {}".format(kbd.get_manufacturer_string()))
            print("Product:      {}".format(kbd.get_product_string()))
        return kbd
    raise ValueError(
        "QMK keyboard not found (VID=0x{:04x} PID=0x{:04x}). "
        "Is it connected and in USB mode?".format(VID, PID)
    )


# ── VIA communication ──────────────────────────────────────────────────────────


def send_via_command(kbd, command_id, channel_id, value_id=0x00, *extra_bytes):
    """Send one VIA command (33 bytes: report-ID 0x00 + 32-byte payload).

    VIA protocol v12 layout: [command_id, channel_id, value_id, value_data...]
    The report-ID byte (buf[0] = 0x00) is stripped by the HID layer before
    the firmware receives the remaining 32 bytes.
    """
    buf = bytearray(REPORT_SIZE + 1)
    buf[0] = 0x00  # report ID
    buf[1] = command_id
    buf[2] = channel_id
    buf[3] = value_id
    for i, b in enumerate(extra_bytes):
        buf[4 + i] = b
    if debug > 1:
        print("HID write: {}".format(buf[: max(6, 4 + len(extra_bytes) + 1)].hex()))
    kbd.write(bytes(buf))
    time.sleep(0.05)


def get_via_value(kbd, channel_id, value_id, response_bytes=1):
    """Query one VIA value; returns a list of response_bytes ints, or None on error."""
    buf = bytearray(REPORT_SIZE + 1)
    buf[0] = 0x00
    buf[1] = VIA_CUSTOM_GET_VALUE
    buf[2] = channel_id
    buf[3] = value_id
    if debug > 1:
        print("HID get query: {}".format(buf[:6].hex()))
    kbd.write(bytes(buf))
    time.sleep(0.05)
    resp = kbd.read(REPORT_SIZE, timeout_ms=500)
    if debug > 1:
        print("HID get resp:  {}".format(bytes(resp).hex() if resp else "(none)"))
    if resp and resp[0] != 0xFF:
        return list(resp[3 : 3 + response_bytes])
    return None


# ── High-level lighting functions ──────────────────────────────────────────────


def lights_off(kbd):
    """Set effect to 0 (RGB_MATRIX_NONE) and save — proper off, not just dim."""
    send_via_command(kbd, VIA_CUSTOM_SET_VALUE, RGB_MATRIX_CHANNEL, QMK_EFFECT, 0x00)
    send_via_command(kbd, VIA_CUSTOM_SAVE, RGB_MATRIX_CHANNEL)
    if debug > 0:
        print("Lights off")


def apply_settings(
    kbd, *, effect=None, brightness=None, hue=None, saturation=None, speed=None
):
    """Send whichever settings are not None, then save to EEPROM.

    If only one of hue/saturation is supplied, the other is read from the
    keyboard so the colour command always sends both bytes.
    """
    if effect is not None:
        send_via_command(
            kbd, VIA_CUSTOM_SET_VALUE, RGB_MATRIX_CHANNEL, QMK_EFFECT, effect
        )
    if brightness is not None:
        send_via_command(
            kbd, VIA_CUSTOM_SET_VALUE, RGB_MATRIX_CHANNEL, QMK_BRIGHTNESS, brightness
        )
    if hue is not None or saturation is not None:
        if hue is None or saturation is None:
            current = get_via_value(
                kbd, RGB_MATRIX_CHANNEL, QMK_COLOR, response_bytes=2
            )
            if current:
                if hue is None:
                    hue = current[0]
                if saturation is None:
                    saturation = current[1]
            else:
                if hue is None:
                    hue = 0
                if saturation is None:
                    saturation = 0
        send_via_command(
            kbd, VIA_CUSTOM_SET_VALUE, RGB_MATRIX_CHANNEL, QMK_COLOR, hue, saturation
        )
    if speed is not None:
        send_via_command(
            kbd, VIA_CUSTOM_SET_VALUE, RGB_MATRIX_CHANNEL, QMK_SPEED, speed
        )
    send_via_command(kbd, VIA_CUSTOM_SAVE, RGB_MATRIX_CHANNEL)
    if debug > 0:
        applied = {
            k: v
            for k, v in dict(
                effect=effect,
                brightness=brightness,
                hue=hue,
                saturation=saturation,
                speed=speed,
            ).items()
            if v is not None
        }
        print("Applied: {}".format(applied))


def lights_on(kbd):
    """Solid white preset (effect=1, brightness=255, saturation=0)."""
    apply_settings(
        kbd,
        effect=ON_EFFECT,
        brightness=ON_BRIGHTNESS,
        hue=ON_HUE,
        saturation=ON_SATURATION,
        speed=ON_SPEED,
    )


def lights_per_key(kbd):
    """Per-key RGB preset — per_key_rgb (effect 23, uses PERKEY_* constants)."""
    apply_settings(
        kbd,
        effect=PERKEY_EFFECT,
        brightness=PERKEY_BRIGHTNESS,
        hue=PERKEY_HUE,
        saturation=PERKEY_SATURATION,
        speed=PERKEY_SPEED,
    )


def lights_programmed(kbd):
    """Programmed preset — reactive_multiwide (effect 19, uses PROGRAMMED_* constants)."""
    apply_settings(
        kbd,
        effect=PROGRAMMED_EFFECT,
        brightness=PROGRAMMED_BRIGHTNESS,
        hue=PROGRAMMED_HUE,
        saturation=PROGRAMMED_SATURATION,
        speed=PROGRAMMED_SPEED,
    )


def lights_query(kbd):
    """Read and display current RGB matrix settings from the keyboard."""
    effect = get_via_value(kbd, RGB_MATRIX_CHANNEL, QMK_EFFECT)
    brightness = get_via_value(kbd, RGB_MATRIX_CHANNEL, QMK_BRIGHTNESS)
    speed = get_via_value(kbd, RGB_MATRIX_CHANNEL, QMK_SPEED)
    color = get_via_value(kbd, RGB_MATRIX_CHANNEL, QMK_COLOR, response_bytes=2)

    if effect:
        eff_name = EFFECT_NAMES.get(effect[0], "(firmware-specific — see --help)")
        print("Effect:     {} — {}".format(effect[0], eff_name))
    else:
        print("Effect:     ?")
    print("Brightness: {}".format(brightness[0] if brightness else "?"))
    if color:
        sat_desc = (
            "white" if color[1] == 0 else "fully saturated" if color[1] == 255 else ""
        )
        sat_str = "{}{}".format(color[1], "  ({})".format(sat_desc) if sat_desc else "")
        print("Hue:        {}".format(color[0]))
        print("Saturation: {}".format(sat_str))
    else:
        print("Hue:        ?")
        print("Saturation: ?")
    print("Speed:      {}".format(speed[0] if speed else "?"))


# ── Screensaver integration ────────────────────────────────────────────────────


def screensaver_message_callback(session, message):
    global debug
    for screensaver in SCREENSAVER_INTERFACES:
        if message.get_interface() != screensaver:
            continue
        if debug > 2:
            pprint.pprint(message)
        if message.get_member() in ("ActiveChanged", "WakeUpScreen"):
            try:
                if debug > 1:
                    print(
                        "{} ActiveChanged/WakeUpScreen: args={}".format(
                            screensaver, message.get_args_list()
                        )
                    )
            except Exception:
                pass
            # Wrap individually so a non-supporting screensaver doesn't
            # block one that does support GetActive.
            try:
                path = "/{}".format(screensaver.replace(".", "/"))
                obj = session.get_object(screensaver, path)
                iface = dbus.Interface(obj, screensaver)
                time.sleep(0.25)
                active = bool(iface.GetActive())
                if debug > 1:
                    print("{} GetActive → {}".format(screensaver, active))
                if active:
                    kbd = find_device()
                    lights_off(kbd)
                    kbd.close()
            except Exception:
                continue

        if message.get_member() == "WakeUpScreen":
            try:
                time.sleep(0.25)
                if debug > 1:
                    print("{} WakeUpScreen → restore lights".format(screensaver))
                kbd = find_device()
                lights_on(kbd)
                kbd.close()
            except Exception:
                continue


# ── Argument parsing and main ──────────────────────────────────────────────────


def byte_value(s):
    """argparse type: accept decimal or 0x hex, validate 0–255."""
    n = int(s, 0)
    if not 0 <= n <= 255:
        raise argparse.ArgumentTypeError("value must be 0–255, got {}".format(n))
    return n


def main():
    global debug

    parser = argparse.ArgumentParser(
        description="Control RGB backlight on Keychron K3 Max (QMK/VIA protocol v12)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=HELP_EPILOG,
    )

    # ── Mode flags (mutually exclusive) ───────────────────────────────────────
    mode_group = parser.add_mutually_exclusive_group()
    mode_group.add_argument(
        "--on",
        "-1",
        dest="mode",
        action="store_const",
        const="on",
        help="Preset: solid white (effect=1, brightness=255, hue=0, saturation=0)",
    )
    mode_group.add_argument(
        "--off",
        "-0",
        dest="mode",
        action="store_const",
        const="off",
        help="Turn backlight off (sets effect=0, saves to EEPROM)",
    )
    mode_group.add_argument(
        "--per-key",
        "-k",
        dest="mode",
        action="store_const",
        const="per_key",
        help="Preset: per_key_rgb (effect 23, uses PERKEY_* constants)",
    )
    mode_group.add_argument(
        "--program",
        "-p",
        dest="mode",
        action="store_const",
        const="program",
        help="Preset: reactive_multiwide (effect 19, uses PROGRAMMED_* constants)",
    )
    mode_group.add_argument(
        "--query",
        "-q",
        dest="mode",
        action="store_const",
        const="query",
        help="Read and display current effect, brightness, hue, saturation and speed",
    )
    mode_group.add_argument(
        "--screensaver",
        "-s",
        dest="mode",
        action="store_const",
        const="screensaver",
        help="Watch screensaver DBus signals: off on lock, solid white on wake",
    )

    # ── Individual settings (can augment a mode or stand alone) ───────────────
    parser.add_argument(
        "--effect",
        "-e",
        type=byte_value,
        metavar="0-255",
        help=(
            "Set effect by ID.  0=off, 1=solid_color; higher IDs are "
            "firmware-specific (see effect table in --help)."
        ),
    )
    parser.add_argument(
        "--brightness",
        "-b",
        type=byte_value,
        metavar="0-255",
        help="LED brightness.  0=off, 128=half, 255=full.",
    )
    parser.add_argument(
        "--hue",
        type=byte_value,
        metavar="0-255",
        help=(
            "Colour hue on the HSV wheel.  "
            "0=red  43=yellow  85=green  128=cyan  170=blue  213=magenta. "
            "Has no visible effect when --saturation is 0 (white)."
        ),
    )
    parser.add_argument(
        "--saturation",
        "--sat",
        type=byte_value,
        metavar="0-255",
        help="Colour saturation.  0=white, 255=fully saturated colour.",
    )
    parser.add_argument(
        "--speed",
        type=byte_value,
        metavar="0-255",
        help="Animation speed.  0=slowest, 128=mid, 255=fastest.",
    )

    parser.add_argument(
        "--verbose",
        "-v",
        action="count",
        default=0,
        help="Increase verbosity (repeat for more detail: -v, -vv, -vvv).",
    )

    args = parser.parse_args()

    individual = dict(
        effect=args.effect,
        brightness=args.brightness,
        hue=args.hue,
        saturation=args.saturation,
        speed=args.speed,
    )
    has_individual = any(v is not None for v in individual.values())

    if args.mode is None and not has_individual:
        parser.print_help()
        sys.exit(1)

    debug = args.verbose

    # ── Screensaver mode: long-running DBus loop ───────────────────────────────
    if args.mode == "screensaver":
        dbus_loop = dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)
        session = dbus.SessionBus(mainloop=dbus_loop)
        for iface in SCREENSAVER_INTERFACES:
            session.add_match_string_non_blocking("interface='{}'".format(iface))
        session.add_message_filter(screensaver_message_callback)
        gi.repository.GLib.MainLoop().run()
        return

    # ── All other modes: open device, act, close ───────────────────────────────
    kbd = find_device()

    if args.mode == "query":
        lights_query(kbd)

    elif args.mode == "off":
        lights_off(kbd)

    else:
        # Start with profile defaults, then override with explicit args.
        settings = {}
        if args.mode == "on":
            settings = dict(
                effect=ON_EFFECT,
                brightness=ON_BRIGHTNESS,
                hue=ON_HUE,
                saturation=ON_SATURATION,
                speed=ON_SPEED,
            )
        elif args.mode == "per_key":
            settings = dict(
                effect=PERKEY_EFFECT,
                brightness=PERKEY_BRIGHTNESS,
                hue=PERKEY_HUE,
                saturation=PERKEY_SATURATION,
                speed=PERKEY_SPEED,
            )
        elif args.mode == "program":
            settings = dict(
                effect=PROGRAMMED_EFFECT,
                brightness=PROGRAMMED_BRIGHTNESS,
                hue=PROGRAMMED_HUE,
                saturation=PROGRAMMED_SATURATION,
                speed=PROGRAMMED_SPEED,
            )
        # Individual args override profile defaults.
        for key, val in individual.items():
            if val is not None:
                settings[key] = val
        apply_settings(kbd, **settings)

    kbd.close()


if __name__ == "__main__":
    main()