"""WS2812 via PWM/DMA — Pi 4 and Zero only (not Pi 5)."""

from __future__ import annotations

try:
    from rpi_ws281x import PixelStrip, ws
except ImportError as exc:  # pragma: no cover - import guard
    PixelStrip = None
    ws = None
    _WS281X_IMPORT_ERROR = exc
else:
    _WS281X_IMPORT_ERROR = None


class Ws281xBackend:
    def __init__(
        self,
        led_count: int,
        *,
        pin: int = 18,
        freq_hz: int = 800_000,
        dma: int = 10,
        invert: bool = False,
        channel: int = 0,
    ) -> None:
        if PixelStrip is None:
            raise ImportError(
                "rpi-ws281x is required for the ws281x backend; install portal requirements"
            ) from _WS281X_IMPORT_ERROR

        self._led_count = led_count
        self._strip = PixelStrip(
            led_count,
            pin,
            freq_hz,
            dma,
            invert,
            255,
            channel,
            strip_type=ws.WS2811_STRIP_GRB,
        )
        self._strip.begin()

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

    def write_grb(self, grb_bytes: bytes) -> None:
        if len(grb_bytes) != self._led_count * 3:
            raise ValueError(
                f"expected {self._led_count * 3} bytes, got {len(grb_bytes)}"
            )

        strip = self._strip
        for i in range(self._led_count):
            base = i * 3
            g, r, b = grb_bytes[base], grb_bytes[base + 1], grb_bytes[base + 2]
            strip.setPixelColor(i, (r << 16) | (g << 8) | b)
        strip.show()

    def clear(self) -> None:
        for i in range(self._led_count):
            self._strip.setPixelColor(i, 0)
        self._strip.show()

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

    def __enter__(self) -> Ws281xBackend:
        return self

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