"""WS2812 over SPI via the rpi5-ws2812 package (Pi 5, GPIO 10)."""

from __future__ import annotations

import numpy as np
from rpi5_ws2812.ws2812 import WS2812SpiDriver

from leds.backends.spi import spidev_bufsize_hint


class Rpi5Ws2812Backend:
    def __init__(
        self,
        led_count: int,
        *,
        spi_bus: int = 0,
        spi_device: int = 0,
        max_speed_hz: int = 4_200_000,
    ) -> None:
        hint = spidev_bufsize_hint(led_count)
        if hint:
            raise OSError(hint)

        self._led_count = led_count
        self._driver = WS2812SpiDriver(spi_bus, spi_device, led_count)
        self._driver._device.max_speed_hz = max_speed_hz

    @property
    def led_count(self) -> int:
        return self._led_count

    def write_grb(self, rgb_bytes: bytes) -> None:
        """RGB bytes passed through to SPI unchanged."""
        if len(rgb_bytes) != self._led_count * 3:
            raise ValueError(
                f"expected {self._led_count * 3} bytes, got {len(rgb_bytes)}"
            )
        buf = np.frombuffer(rgb_bytes, dtype=np.uint8).reshape(self._led_count, 3)
        self._driver.write(buf)

    def clear(self) -> None:
        self._driver.clear()

    def close(self) -> None:
        try:
            self.clear()
        finally:
            self._driver._device.close()

    def __enter__(self) -> Rpi5Ws2812Backend:
        return self

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