"""WS2812 over SPI — works on Pi 5, 4, and Zero when SPI is enabled."""

from __future__ import annotations

from array import array
from pathlib import Path

try:
    from spidev import SpiDev
except ImportError as exc:  # pragma: no cover - import guard
    SpiDev = None
    _SPI_IMPORT_ERROR = exc
else:
    _SPI_IMPORT_ERROR = None

# WS2812 bit patterns via SPI, from https://github.com/mattaw/ws2812_spi_python
_LED_ZERO = 0b1100_0000
_LED_ONE = 0b1111_1100
_PREAMBLE = 42
_DEFAULT_BUFSIZ = 65536


def spidev_bufsize() -> int:
    try:
        return int(Path("/sys/module/spidev/parameters/bufsiz").read_text().strip())
    except (OSError, ValueError):
        return 4096


def spi_frame_bytes(led_count: int) -> int:
    return _PREAMBLE + led_count * 24


def leds_from_truncated_frame(bufsize: int) -> int:
    payload = max(0, bufsize - _PREAMBLE)
    return payload // 24


def spidev_bufsize_hint(led_count: int) -> str:
    required = spi_frame_bytes(led_count)
    bufsize = spidev_bufsize()
    if bufsize >= required:
        return ""
    suggested = max(_DEFAULT_BUFSIZ, required + 1024)
    truncated = leds_from_truncated_frame(bufsize)
    return (
        f"spidev bufsiz is {bufsize} but {led_count} LEDs need a {required}-byte SPI frame "
        f"(only ~{truncated} LEDs update at the current limit).\n"
        f"Fix now:\n"
        f"  sudo modprobe -r spidev\n"
        f"  sudo modprobe spidev bufsiz={suggested}\n"
        f"Permanent:\n"
        f"  pipenv run python examples/setup_spi_bufsiz.py --install"
    )


def _write_spi_frame(device: SpiDev, frame: array) -> None:
    required = len(frame)
    if required > spidev_bufsize():
        raise OSError(spidev_bufsize_hint((required - _PREAMBLE) // 24))
    device.writebytes2(frame)


class SpiBackend:
    def __init__(
        self,
        led_count: int,
        *,
        spi_bus: int = 0,
        spi_device: int = 0,
        max_speed_hz: int = 6_500_000,
    ) -> None:
        if SpiDev is None:
            raise ImportError(
                "spidev is required for the SPI backend; install portal requirements"
            ) from _SPI_IMPORT_ERROR

        self._led_count = led_count
        self._device = SpiDev()
        self._device.open(spi_bus, spi_device)
        self._device.max_speed_hz = max_speed_hz
        self._device.mode = 0b00
        self._device.lsbfirst = False

        frame_len = spi_frame_bytes(led_count)
        self._frame = array("B", [0] * frame_len)
        self._clear_frame = array("B", [0] * frame_len)
        for i in range(_PREAMBLE, frame_len):
            self._clear_frame[i] = _LED_ZERO

        hint = spidev_bufsize_hint(led_count)
        if hint:
            raise OSError(hint)

    @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)}"
            )

        out = self._frame
        pos = _PREAMBLE
        for value in grb_bytes:
            for bit in (128, 64, 32, 16, 8, 4, 2, 1):
                out[pos] = _LED_ONE if value & bit else _LED_ZERO
                pos += 1

        _write_spi_frame(self._device, out)

    def clear(self) -> None:
        _write_spi_frame(self._device, self._clear_frame)

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

    def __enter__(self) -> SpiBackend:
        return self

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