controller_api.py

#!/usr/bin/env python3
"""
DreamScaler USB Controller API
Python library for communicating with the Arduino USB controller for SK6812 LED strips.

Usage:
    from controller_api import LEDController

    controller = LEDController('/dev/ttyACM0')  # or 'COM3' on Windows
    controller.connect()

    # Set a single LED
    controller.set_pixel(0, r=255, g=0, b=0, w=0)
    controller.sync()

    # Turn off all LEDs
    controller.clear_all()

    controller.disconnect()
"""

import serial
import time
import struct
from typing import Tuple, List, Optional
from enum import IntEnum

# Import LED_COUNT from the appropriate map file (uncomment the one you need)
from arturia_keylab49_map import LED_COUNT   # Arturia KeyLab 49 MKII
# from arturia_keylab61_map import LED_COUNT  # Arturia KeyLab 61 MKII


class Command(IntEnum):
    """Protocol commands."""
    # System
    PING = 0x01
    GET_INFO = 0x02
    RESET = 0x03
    
    # LED configuration
    SET_LED_COUNT = 0x10
    SET_LED_PIN = 0x11
    INIT_STRIP = 0x12
    RESIZE_STRIP = 0x13
    
    # LED control – single pixel
    SET_PIXEL_RGBW = 0x20
    SET_PIXEL_RGB = 0x21
    SET_PIXEL_W = 0x22
    
    # LED control – bulk
    SET_RANGE_RGBW = 0x30
    SET_ALL_RGBW = 0x31
    CLEAR_ALL = 0x32
    
    # Buffer mode
    BUFFER_START = 0x40
    BUFFER_PIXEL = 0x41
    BUFFER_END = 0x42
    
    # Stream mode
    STREAM_START = 0x50
    STREAM_DATA = 0x51
    STREAM_END = 0x52
    
    # Bulk update (optimised)
    BULK_UPDATE = 0x55
    
    # Synchronisation
    SYNC = 0x60

    # Effects
    FILL_GRADIENT = 0x70
    BRIGHTNESS = 0x71


class Response(IntEnum):
    """Responses from the controller."""
    OK = 0xF0
    PONG = 0xF1
    INFO = 0xF2
    ERROR = 0xFE
    UNKNOWN_CMD = 0xFF


class ErrorCode(IntEnum):
    """Error codes."""
    BUFFER_OVERFLOW = 0x01
    INVALID_PARAM = 0x02
    NOT_INITIALIZED = 0x03
    OUT_OF_RANGE = 0x04
    TIMEOUT = 0x05


class LEDControllerError(Exception):
    """Exception raised for controller errors."""
    pass


class LEDController:
    """API for controlling an SK6812 LED strip via USB."""
    
    def __init__(self, port: str, baud_rate: int = 115200, timeout: float = 2.0):
        """
        Initialise the controller.

        Args:
            port: Serial port (e.g. '/dev/ttyACM0' or 'COM3')
            baud_rate: Communication speed (default 115200)
            timeout: Read timeout in seconds
        """
        self.port = port
        self.baud_rate = baud_rate
        self.timeout = timeout
        self.serial: Optional[serial.Serial] = None
        self.led_count = LED_COUNT  # defined in arturia_keylab49_map.py
        self.led_pin = 6
        self.is_initialized = False
    
    def connect(self) -> bool:
        """
        Connect to the controller.

        Returns:
            True if successful.
        """
        try:
            self.serial = serial.Serial(
                self.port,
                self.baud_rate,
                timeout=self.timeout
            )
            # Wait for Arduino to initialise
            time.sleep(2)

            # Flush any data from the buffer (Arduino bootloader, init messages)
            self.serial.reset_input_buffer()
            self.serial.reset_output_buffer()

            # Test connection
            if self.ping():
                info = self.get_info()
                if info:
                    self.led_count = info['led_count']
                    self.led_pin = info['led_pin']
                    self.is_initialized = info['initialized']
                    print(f'Connected to controller: {info}')
                    return True

            return False
        except Exception as e:
            print(f'Connection error: {e}')
            return False
    
    def disconnect(self):
        """
        Disconnect with aggressive port release.

        IMPORTANT: Implements the "baudrate trick" to properly release the port.
        After program exit (including Ctrl+C) the port can remain locked in the OS.
        The fix is to re-open the port at a different baud rate and immediately
        close it — this resets the port state in the OS and frees it.

        Without this step the port stays locked after Ctrl+C and you would need
        to run release_port.py or physically disconnect USB.
        """
        if self.serial:
            port_name = self.port  # save port name for baudrate trick
            
            try:
                if self.serial.is_open:
                    # Graceful close attempt
                    try:
                        self.serial.reset_input_buffer()
                        self.serial.reset_output_buffer()
                    except:
                        pass

                    # Set DTR and RTS to False (may help release the port)
                    try:
                        self.serial.dtr = False
                        self.serial.rts = False
                    except:
                        pass

                    time.sleep(0.1)

                    # Close port
                    try:
                        self.serial.close()
                    except:
                        pass

            except Exception as e:
                pass  # Ignore errors
            finally:
                self.serial = None
                self.is_initialized = False
            
            # ============================================================
            # BAUDRATE TRICK — KEY STEP FOR RELEASING THE PORT
            # ============================================================
            # When the port freezes at 115200, opening it at a different baud
            # (9600) and immediately closing resets the port state in Windows.
            # Without this the port stays locked after Ctrl+C!
            try:
                time.sleep(0.1)
                temp_serial = serial.Serial(port_name, baudrate=9600, timeout=0.1)
                temp_serial.reset_input_buffer()
                temp_serial.reset_output_buffer()
                temp_serial.dtr = False
                temp_serial.rts = False
                time.sleep(0.1)
                temp_serial.close()
                time.sleep(0.2)  # Important: wait for the OS to release the port
            except:
                pass  # If it fails, no problem
    
    def _send_command(self, cmd: int, data: bytes = b'') -> bytes:
        """
        Send a command and wait for the response.

        Args:
            cmd: Command byte
            data: Additional data bytes

        Returns:
            Response from the controller.
        """
        if not self.serial or not self.serial.is_open:
            raise LEDControllerError('Controller not connected')

        # Send the command
        self.serial.write(bytes([cmd]) + data)
        self.serial.flush()

        # Read the response
        response = self.serial.read(1)
        if len(response) == 0:
            raise LEDControllerError('Timeout — no response')

        resp_code = response[0]

        if resp_code == Response.ERROR:
            error_code = self.serial.read(1)
            if len(error_code) > 0:
                raise LEDControllerError(f'Controller error: {ErrorCode(error_code[0]).name}')
            else:
                raise LEDControllerError('Unknown controller error')

        return bytes([resp_code])
    
    def _wait_for_ok(self):
        """Wait for an OK response."""
        resp = self._send_command(0, b'')  # dummy — response is read inside _send_command
        # The actual response is retrieved in _send_command
    
    # ========== SYSTEM COMMANDS ==========

    def ping(self) -> bool:
        """Test the connection."""
        try:
            resp = self._send_command(Command.PING)
            return resp[0] == Response.PONG
        except:
            return False
    
    def get_info(self) -> Optional[dict]:
        """
        Read controller information.

        Returns:
            Dict with info, or None on failure.
        """
        try:
            # Flush the buffer
            self.serial.reset_input_buffer()

            # Send command
            self.serial.write(bytes([Command.GET_INFO]))
            self.serial.flush()

            # Read response
            resp = self.serial.read(1)
            
            if len(resp) > 0 and resp[0] == Response.INFO:
                data = self.serial.read(6)  # Protocol version + 5 bytes
                if len(data) == 6:
                    return {
                        'protocol_version': data[0],
                        'led_count': (data[1] << 8) | data[2],
                        'led_pin': data[3],
                        'initialized': bool(data[4]),
                        'brightness': data[5]
                    }
            return None
        except Exception as e:
            print(f'Error reading info: {e}')
            return None

    def reset(self):
        """Reset the controller."""
        self._send_command(Command.RESET)
    
    # ========== CONFIGURATION ==========

    def set_led_count(self, count: int):
        """Set the number of LEDs (does NOT reinitialise the strip — call init_strip() afterwards)."""
        data = struct.pack('>H', count)  # Big-endian uint16
        self._send_command(Command.SET_LED_COUNT, data)
        self.led_count = count

    def resize_strip(self, count: int):
        """
        Set the number of LEDs and immediately reinitialise the strip.

        This is the recommended way to change strip length at runtime — it sends
        a single CMD_RESIZE_STRIP command which atomically updates ledCount and
        calls initializeStrip() on the firmware side.

        Args:
            count: New number of LEDs (1-LED_COUNT)
        """
        if not (1 <= count <= LED_COUNT):
            raise ValueError(f'LED count must be between 1 and {LED_COUNT}, got {count}')
        data = struct.pack('>H', count)  # Big-endian uint16
        self._send_command(Command.RESIZE_STRIP, data)
        self.led_count = count
        self.is_initialized = True

    def set_led_pin(self, pin: int):
        """Set the output pin."""
        data = bytes([pin])
        self._send_command(Command.SET_LED_PIN, data)
        self.led_pin = pin
    
    def init_strip(self):
        """Initialise the LED strip with the current settings."""
        self._send_command(Command.INIT_STRIP)
        self.is_initialized = True
    
    # ========== SINGLE-PIXEL CONTROL ==========

    def set_pixel(self, index: int, r: int = 0, g: int = 0, b: int = 0, w: int = 0):
        """
        Set a single LED colour (RGBW).

        Args:
            index: LED index (0-based)
            r, g, b, w: Colour values 0-255
        """
        data = struct.pack('>HBBBB', index, r, g, b, w)
        self._send_command(Command.SET_PIXEL_RGBW, data)
    
    def set_pixel_rgb(self, index: int, r: int, g: int, b: int):
        """Set an LED colour with RGB only (W=0)."""
        data = struct.pack('>HBBB', index, r, g, b)
        self._send_command(Command.SET_PIXEL_RGB, data)
    
    def set_pixel_white(self, index: int, w: int):
        """Set an LED using the White channel only."""
        data = struct.pack('>HB', index, w)
        self._send_command(Command.SET_PIXEL_W, data)
    
    # ========== BULK OPERATIONS ==========

    def set_range(self, start: int, end: int, r: int = 0, g: int = 0, b: int = 0, w: int = 0):
        """
        Set a range of LEDs to the same colour.

        Args:
            start: Start index
            end: End index (inclusive)
            r, g, b, w: Colour values 0-255
        """
        data = struct.pack('>HHBBBB', start, end, r, g, b, w)
        self._send_command(Command.SET_RANGE_RGBW, data)
    
    def set_all(self, r: int = 0, g: int = 0, b: int = 0, w: int = 0):
        """Set all LEDs to the same colour."""
        data = struct.pack('BBBB', r, g, b, w)
        self._send_command(Command.SET_ALL_RGBW, data)
    
    def clear_all(self):
        """Turn off all LEDs."""
        self._send_command(Command.CLEAR_ALL)
    
    # ========== BUFFER MODE ==========

    def buffer_begin(self):
        """Begin a buffered update."""
        self._send_command(Command.BUFFER_START)
    
    def buffer_set_pixel(self, index: int, r: int = 0, g: int = 0, b: int = 0, w: int = 0):
        """Set an LED in buffer mode."""
        data = struct.pack('>HBBBB', index, r, g, b, w)
        self._send_command(Command.BUFFER_PIXEL, data)
    
    def buffer_end(self):
        """End buffered update and sync to the strip."""
        self._send_command(Command.BUFFER_END)
    
    def buffer_update(self, pixels: List[Tuple[int, int, int, int, int]]):
        """
        Batch-update LEDs using buffer mode.

        Args:
            pixels: List of (index, r, g, b, w) tuples
        """
        self.buffer_begin()
        for pixel in pixels:
            self.buffer_set_pixel(*pixel)
        self.buffer_end()
    
    # ========== STREAM MODE (FASTEST) ==========

    def stream_begin(self, count: int):
        """Begin a stream update."""
        data = struct.pack('>H', count)
        self._send_command(Command.STREAM_START, data)
    
    def stream_pixel(self, r: int, g: int, b: int, w: int):
        """Send one LED in stream mode (sequential)."""
        data = struct.pack('BBBB', r, g, b, w)
        self._send_command(Command.STREAM_DATA, data)
    
    def stream_end(self):
        """End stream and sync to the strip."""
        self._send_command(Command.STREAM_END)
    
    def stream_update(self, pixels: List[Tuple[int, int, int, int]], wait_response: bool = False):
        """
        Fastest way to update all LEDs.
        Uses the BULK_UPDATE command for minimal overhead.

        Args:
            pixels: List of (r, g, b, w) tuples starting from index 0
            wait_response: Ignored (kept for compatibility)
        """
        if not self.serial or not self.serial.is_open:
            raise LEDControllerError('Controller not connected')

        # Flush the buffer
        self.serial.reset_input_buffer()

        # Build the BULK_UPDATE packet: CMD + count(2B) + all RGBW data
        packet = bytearray()
        packet.append(Command.BULK_UPDATE)
        packet.extend(struct.pack('>H', len(pixels)))  # count

        # Send the command and count
        self.serial.write(packet)
        self.serial.flush()

        # Wait for OK — Arduino is ready (longer timeout for large payloads)
        self.serial.timeout = 5.0
        resp = self.serial.read(1)
        if len(resp) == 0 or resp[0] != Response.OK:
            raise LEDControllerError('Arduino not ready for bulk update')

        # Send all RGBW data at once
        data = bytearray()
        for pixel in pixels:
            data.extend(struct.pack('BBBB', *pixel))

        self.serial.write(data)
        self.serial.flush()

        # Wait for final OK (after sync) — longer timeout
        self.serial.timeout = 5.0
        resp = self.serial.read(1)

        # Reset timeout
        self.serial.timeout = self.timeout

        if len(resp) == 0:
            raise LEDControllerError('Bulk update timeout — no response')

        if resp[0] == Response.ERROR:
            # Read error code
            error_code = self.serial.read(1)
            if len(error_code) > 0:
                error_names = {0x01: 'BUFFER_OVERFLOW', 0x02: 'INVALID_PARAM', 
                              0x03: 'NOT_INITIALIZED', 0x04: 'OUT_OF_RANGE', 0x05: 'TIMEOUT'}
                error_name = error_names.get(error_code[0], f'UNKNOWN({error_code[0]:02x})')
                raise LEDControllerError(f"Arduino error: {error_name}")
            else:
                raise LEDControllerError("Arduino error (no code)")
        
        if resp[0] != Response.OK:
            raise LEDControllerError(f"Bulk update unexpected response: {resp[0]:02x}")
    
    # ========== SYNCHRONISATION ==========

    def sync(self):
        """Immediately sync — push buffered data to the LED strip."""
        self._send_command(Command.SYNC)
    
    # ========== EFFECTS ==========

    def fill_gradient(self, start: int, end: int,
                     r1: int, g1: int, b1: int, w1: int,
                     r2: int, g2: int, b2: int, w2: int):
        """
        Fill a gradient between two colours.

        Args:
            start, end: LED range
            r1, g1, b1, w1: Start colour
            r2, g2, b2, w2: End colour
        """
        data = struct.pack('>HHBBBBBBBB', start, end, r1, g1, b1, w1, r2, g2, b2, w2)
        self._send_command(Command.FILL_GRADIENT, data)
    
    def set_brightness(self, brightness: int):
        """
        Set global brightness.

        Args:
            brightness: 0-255 (0=off, 255=full)
        """
        data = bytes([brightness])
        self._send_command(Command.BRIGHTNESS, data)
    
    # ========== UTILITY ==========

    def rgb_to_rgbw(self, r: int, g: int, b: int) -> Tuple[int, int, int, int]:
        """
        Convert RGB to RGBW (extract white component).

        Args:
            r, g, b: RGB values 0-255

        Returns:
            Tuple (r, g, b, w)
        """
        # Simple method: white = minimum of RGB
        w = min(r, g, b)
        return (r - w, g - w, b - w, w)
    
    def hsv_to_rgb(self, h: float, s: float, v: float) -> Tuple[int, int, int]:
        """
        Convert HSV to RGB.

        Args:
            h: Hue 0-360
            s: Saturation 0-1
            v: Value 0-1

        Returns:
            Tuple (r, g, b) 0-255
        """
        import colorsys
        r, g, b = colorsys.hsv_to_rgb(h/360, s, v)
        return (int(r * 255), int(g * 255), int(b * 255))
    
    def __enter__(self):
        """Context manager support"""
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager cleanup"""
        try:
            # Turn off all LEDs before disconnecting
            if self.serial and self.serial.is_open:
                try:
                    self.clear_all()
                except:
                    pass  # Ignore errors during final cleanup
        except:
            pass
        finally:
            self.disconnect()
        return False  # Do not suppress exceptions


# ========== USAGE EXAMPLES ==========

def example_basic():
    """Basic usage example."""
    with LEDController('/dev/ttyACM0') as controller:
        # Red LED at index 0
        controller.set_pixel(0, r=255, g=0, b=0, w=0)
        controller.sync()
        time.sleep(1)

        # Green LED at index 1
        controller.set_pixel(1, r=0, g=255, b=0, w=0)
        controller.sync()
        time.sleep(1)

        # All white
        controller.set_all(r=0, g=0, b=0, w=50)
        time.sleep(1)

        # Clear
        controller.clear_all()


def example_gradient():
    """Gradient example."""
    with LEDController('/dev/ttyACM0') as controller:
        # Gradient from red to blue
        controller.fill_gradient(
            0, 143,
            255, 0, 0, 0,    # Red
            0, 0, 255, 0     # Blue
        )
        controller.sync()


def example_animation():
    """Animation example."""
    with LEDController('/dev/ttyACM0') as controller:
        # Rainbow animation
        for offset in range(360):
            pixels = []
            for i in range(LED_COUNT):
                hue = (i * 360 / LED_COUNT + offset) % 360
                r, g, b = controller.hsv_to_rgb(hue, 1.0, 0.5)
                pixels.append((r, g, b, 0))
            controller.stream_update(pixels)
            time.sleep(0.02)


if __name__ == '__main__':
    print('DreamScaler USB Controller API')
    print('Import with: from controller_api import LEDController')
    print('\nRunning test animation...')
    
    # Run an example
    # example_basic()
    # example_gradient()
    # example_animation()