"""Per-key RGB keyboard (Chicony MCU 04F2:0117/011A, interface 3, vendor page FF02) via hidraw.

Protocol (from PredatorSense USB captures, see the Venator project):
  * 8-byte feature reports  <op> p1..p6 <checksum>,  checksum = 0xFF - (sum & 0xFF)
  * 512-byte output report  carrying the per-key framebuffer (128 cells x {0,R,G,B})
"""

import fcntl
import glob
import os
import threading

ACER_VID = 0x04F2
KNOWN_PIDS = (0x0117, 0x011A)
LED_INTERFACE = 3
NUM_CELLS = 128

OP_BEGIN = 0x88
OP_MODE_ZONE = 0xB1
OP_MODE_PERKEY = 0x12
OP_SET_COLOR = 0x14
OP_APPLY = 0x08
SUB_APPLY = 0x02
MODE_TAG = 0x05
FLAG_DEFAULT = 0x01
SCOPE_ZONE = 0x01
SCOPE_PERKEY = 0x08

# name -> (EFF byte, uses colour?, display name)
EFFECTS = {
    "static":    (0x01, True,  "Static"),
    "breathing": (0x02, True,  "Breathing"),
    "rainbow":   (0x03, False, "Rainbow wave"),
    "snake":     (0x05, True,  "Snake"),
    "ripple":    (0x06, True,  "Ripple (reactive)"),
    "neon":      (0x08, False, "Neon"),
    "rain":      (0x0A, True,  "Rain"),
    "explosion": (0x12, True,  "Explosion"),
    "pulse":     (0x25, True,  "Pulse (reactive)"),
    "stars":     (0x26, True,  "Stars"),
    "meteor":    (0x27, True,  "Meteor (reactive)"),
    "aura":      (0x28, True,  "Aura (reactive)"),
    "perkey":    (0x33, False, "Custom (per key)"),
    "off":       (0x01, False, "Off"),
}


def _ioc(direction, type_, nr, size):
    return (direction << 30) | (size << 16) | (ord(type_) << 8) | nr


def HIDIOCSFEATURE(length):
    return _ioc(3, "H", 0x06, length)  # _IOC_WRITE|_IOC_READ


def checksum(body7):
    return (0xFF - (sum(body7) & 0xFF)) & 0xFF


def find_hidraw():
    """Path of the hidraw node of the LED interface, or None."""
    for node in sorted(glob.glob("/sys/class/hidraw/hidraw*")):
        try:
            with open(os.path.join(node, "device", "uevent")) as f:
                uevent = f.read()
        except OSError:
            continue
        hid_id = next((l.split("=", 1)[1] for l in uevent.splitlines()
                       if l.startswith("HID_ID=")), "")
        parts = hid_id.split(":")
        if len(parts) != 3:
            continue
        vid, pid = int(parts[1], 16), int(parts[2], 16)
        if vid != ACER_VID or pid not in KNOWN_PIDS:
            continue
        devpath = os.path.realpath(os.path.join(node, "device"))
        # .../1-7:1.3/0003:04F2:011A.0004  -> interface number from the parent
        iface = os.path.basename(os.path.dirname(devpath)).rsplit(".", 1)[-1]
        if iface == str(LED_INTERFACE):
            return "/dev/" + os.path.basename(node)
    return None


class Keyboard:
    def __init__(self):
        self._lock = threading.Lock()

    @property
    def path(self):
        return find_hidraw()

    def available(self):
        return self.path is not None

    def _cmd(self, fd, op, p1=0, p2=0, p3=0, p4=0, p5=0, p6=0):
        body = [op, p1, p2, p3, p4, p5, p6]
        buf = bytearray([0] + body + [checksum(body)])
        fcntl.ioctl(fd, HIDIOCSFEATURE(len(buf)), buf)

    def apply(self, effect="static", color=(255, 255, 255), brightness=200,
              frame=None, effect_id=0):
        """Apply a hardware effect or (effect='perkey') a per-key image.

        frame: bytes/list with 128*3 RGB values (perkey only).
        effect_id: optional raw EFF byte (effect variants with other speed/direction).
        """
        path = self.path
        if not path:
            raise RuntimeError("RGB keyboard (hidraw) not found")
        brightness = max(0, min(255, int(brightness)))
        r, g, b = (max(0, min(255, int(c))) for c in color)
        with self._lock:
            fd = os.open(path, os.O_RDWR)
            try:
                self._cmd(fd, OP_BEGIN)
                if effect == "perkey":
                    data = bytes(frame or bytes(NUM_CELLS * 3))
                    data = data[:NUM_CELLS * 3].ljust(NUM_CELLS * 3, b"\0")
                    self._cmd(fd, OP_MODE_PERKEY, 0, 0, 8, 0, 0, 0)
                    out = bytearray(1 + NUM_CELLS * 4)
                    for i in range(NUM_CELLS):
                        out[1 + i * 4 + 1:1 + i * 4 + 4] = data[i * 3:i * 3 + 3]
                    os.write(fd, bytes(out))
                    self._cmd(fd, OP_APPLY, SUB_APPLY, 0x33, MODE_TAG,
                              brightness, SCOPE_PERKEY, FLAG_DEFAULT)
                    return
                self._cmd(fd, OP_MODE_ZONE)
                if effect == "off" or brightness == 0:
                    self._cmd(fd, OP_APPLY, SUB_APPLY, 0x01, MODE_TAG, 0,
                              SCOPE_ZONE, FLAG_DEFAULT)
                    return
                if effect not in EFFECTS:
                    raise ValueError("Unknown effect: %s" % effect)
                self._cmd(fd, OP_SET_COLOR, 0, 0, r, g, b, 0)
                eff = effect_id or EFFECTS[effect][0]
                self._cmd(fd, OP_APPLY, SUB_APPLY, eff, MODE_TAG, brightness,
                          SCOPE_ZONE, FLAG_DEFAULT)
            finally:
                os.close(fd)