portal/leds/strip.py

"""Unified WS2812 strip API for Raspberry Pi 5, 4, and Zero."""

from __future__ import annotations

from typing import Literal, Sequence, Tuple, Union

from leds.backends.pio import PioBackend
from leds.backends.rpi5_ws2812 import Rpi5Ws2812Backend
from leds.backends.spi import SpiBackend
from leds.backends.ws281x import Ws281xBackend
from leds.config import pi5_setup_hint
from leds.detect import board_family, board_model, pio_available, supports_ws281x
from leds.wire import normalize_wire_order, rgb_bytes_to_wire, swap_rg_bytes

Color = Union[Tuple[int, int, int], Sequence[int]]
BackendName = Literal["auto", "pio", "spi", "rpi5_ws2812", "ws281x"]


def _clamp(value: int) -> int:
    return max(0, min(255, int(value)))


def _normalize_color(color: Color) -> Tuple[int, int, int]:
    if len(color) != 3:
        raise ValueError("color must be (r, g, b)")
    return _clamp(color[0]), _clamp(color[1]), _clamp(color[2])


def _apply_brightness(color: Tuple[int, int, int], brightness: float) -> Tuple[int, int, int]:
    scale = max(0.0, min(1.0, brightness))
    return (
        int(color[0] * scale),
        int(color[1] * scale),
        int(color[2] * scale),
    )


def _resolve_backend(name: BackendName) -> str:
    if name != "auto":
        if name == "ws281x" and not supports_ws281x():
            raise RuntimeError(
                f"{board_model()} does not support ws281x; use backend='pio' or 'spi'"
            )
        if name == "pio" and board_family() == "pi5" and not pio_available():
            raise RuntimeError(pi5_setup_hint(10, 10))
        if name == "rpi5_ws2812" and board_family() != "pi5":
            raise RuntimeError("rpi5_ws2812 backend requires Raspberry Pi 5")
        return name

    if board_family() == "pi5":
        return "rpi5_ws2812"
    return "ws281x"


class LedStrip:
    """
    Drive a WS2812 / NeoPixel strip on Pi 5, 4, or Zero.

    Default backends:
      - Pi 5: rpi5-ws2812 SPI on GPIO 10 (/dev/spidev0.0)
      - Pi 5 (pio): ws2812-pio on GPIO 18/13 (/dev/leds0, /dev/leds1) when backend='pio'
      - Pi 4 / Zero: PWM on GPIO 18/13 via rpi_ws281x
    """

    def __init__(
        self,
        count: int,
        *,
        backend: BackendName = "auto",
        pin: int = 18,
        channel: int = 0,
        device: str = "/dev/leds0",
        brightness: float = 1.0,
        spi_bus: int = 0,
        spi_device: int = 0,
        spi_max_speed_hz: int | None = None,
        wire_order: str | None = None,
        passthrough: bool = False,
        swap_rg: bool = True,
    ) -> None:
        if count < 1:
            raise ValueError("count must be at least 1")

        self._count = count
        self._brightness = max(0.0, min(1.0, brightness))
        self._passthrough = passthrough
        self._swap_rg = swap_rg
        self._wire_order = normalize_wire_order(wire_order)
        self._pixels: list[Tuple[int, int, int]] = [(0, 0, 0)] * count

        resolved = _resolve_backend(backend)
        if resolved == "pio":
            self._backend = PioBackend(count, device=device)
        elif resolved == "spi":
            spi_kwargs: dict[str, int] = {}
            if spi_max_speed_hz is not None:
                spi_kwargs["max_speed_hz"] = spi_max_speed_hz
            self._backend = SpiBackend(
                count, spi_bus=spi_bus, spi_device=spi_device, **spi_kwargs
            )
        elif resolved == "rpi5_ws2812":
            hz = spi_max_speed_hz if spi_max_speed_hz is not None else 4_200_000
            self._backend = Rpi5Ws2812Backend(
                count, spi_bus=spi_bus, spi_device=spi_device, max_speed_hz=hz
            )
        else:
            self._backend = Ws281xBackend(count, pin=pin, channel=channel)

        self._backend_name = resolved

    @property
    def count(self) -> int:
        return self._count

    @property
    def backend(self) -> str:
        return self._backend_name

    @property
    def brightness(self) -> float:
        return self._brightness

    @brightness.setter
    def brightness(self, value: float) -> None:
        self._brightness = max(0.0, min(1.0, value))

    def __len__(self) -> int:
        return self._count

    def __getitem__(self, index: int) -> Tuple[int, int, int]:
        return self._pixels[index]

    def __setitem__(self, index: int, color: Color) -> None:
        self._pixels[index] = _normalize_color(color)

    def fill(self, color: Color) -> None:
        rgb = _normalize_color(color)
        self._pixels = [rgb] * self._count

    def set_pixel(self, index: int, color: Color) -> None:
        self[index] = color

    def show(self) -> None:
        self._write_frame(self._logical_frame_bytes())

    def show_rgb_bytes(self, rgb: bytes) -> None:
        """Push pre-dimmed RGB frame to the strip (skips pixel buffer)."""
        if len(rgb) != self._count * 3:
            raise ValueError(f"expected {self._count * 3} bytes, got {len(rgb)}")
        scale = self._brightness
        if scale >= 1.0:
            logical = rgb
        else:
            dimmed = bytearray(len(rgb))
            for i in range(0, len(rgb), 3):
                dimmed[i] = int(rgb[i] * scale)
                dimmed[i + 1] = int(rgb[i + 1] * scale)
                dimmed[i + 2] = int(rgb[i + 2] * scale)
            logical = bytes(dimmed)
        self._write_frame(logical)

    def _write_frame(self, logical_rgb: bytes) -> None:
        frame = swap_rg_bytes(logical_rgb) if self._swap_rg else logical_rgb
        if self._passthrough:
            self._backend.write_grb(frame)
            return
        self._backend.write_grb(self._encode_frame_bytes(frame))

    def _logical_frame_bytes(self) -> bytes:
        logical = bytearray(self._count * 3)
        for i, (r, g, b) in enumerate(self._pixels):
            r, g, b = _apply_brightness((r, g, b), self._brightness)
            base = i * 3
            logical[base] = r
            logical[base + 1] = g
            logical[base + 2] = b
        return bytes(logical)

    def _encode_frame_bytes(self, logical_rgb: bytes) -> bytes:
        return rgb_bytes_to_wire(logical_rgb, self._wire_order)

    def clear(self) -> None:
        self.fill((0, 0, 0))
        self._backend.clear()

    def close(self) -> None:
        self._backend.close()

    def __enter__(self) -> LedStrip:
        return self

    def __exit__(self, *_exc) -> None:
        self.close()