portal/leds/panel_bridge.py

"""UDP-to-SPI bridge — makes a Pi-connected panel look like Pico firmware."""

from __future__ import annotations

import argparse
import queue
import socket
import sys
import threading
from typing import Sequence

from leds.array_config import MATRIX_BRIGHTNESS, panel_layout, panel_pixel_count
from leds.config import panel_strip_config
from leds.matrix import LedMatrix
from leds.panel_udp import DEFAULT_PORT

# Larger kernel RX buffer so SPI show() does not drop bursts from the animator.
_SOCKET_RCVBUF = 1 << 20
_DEFAULT_SPI_MHZ = 4.2


def decode_panel_frame(
    data: bytes,
    pixel_count: int,
    panel_id: int,
) -> bytes | None:
    """Match Pico firmware: raw RGB, or panel_id byte + RGB."""
    rgb_len = pixel_count * 3
    if len(data) == rgb_len:
        return data
    if len(data) == rgb_len + 1:
        prefix = data[0]
        if prefix != 255 and prefix != (panel_id & 0xFF):
            return None
        return data[1:]
    return None


def _recv_latest_frame(
    sock: socket.socket,
    bufsize: int,
    pixel_count: int,
    panel_id: int,
) -> bytes | None:
    """Block for one datagram, then drain the queue and return the newest frame."""
    rgb: bytes | None = None
    while rgb is None:
        try:
            data, _addr = sock.recvfrom(bufsize)
        except InterruptedError:
            continue
        rgb = decode_panel_frame(data, pixel_count, panel_id)
    sock.setblocking(False)
    try:
        while True:
            try:
                data, _addr = sock.recvfrom(bufsize)
            except BlockingIOError:
                break
            newer = decode_panel_frame(data, pixel_count, panel_id)
            if newer is not None:
                rgb = newer
    finally:
        sock.setblocking(True)
    return rgb


def _receiver_loop(
    sock: socket.socket,
    bufsize: int,
    pixel_count: int,
    panel_id: int,
    frames: queue.Queue[bytes],
    stop: threading.Event,
) -> None:
    while not stop.is_set():
        try:
            rgb = _recv_latest_frame(sock, bufsize, pixel_count, panel_id)
        except OSError:
            if stop.is_set():
                return
            raise
        if rgb is None:
            continue
        try:
            frames.put_nowait(rgb)
        except queue.Full:
            try:
                frames.get_nowait()
            except queue.Empty:
                pass
            frames.put_nowait(rgb)


def add_bridge_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument(
        "--panel-index",
        type=int,
        required=True,
        choices=range(5),
        metavar="N",
        help="Panel index (sets LED count / SPI layout)",
    )
    parser.add_argument(
        "--bind",
        default="0.0.0.0",
        help="Bind address (default: 0.0.0.0; animators send to 127.0.0.1)",
    )
    parser.add_argument("--port", type=int, default=DEFAULT_PORT)
    parser.add_argument("--brightness", type=float, default=MATRIX_BRIGHTNESS)
    parser.add_argument(
        "--spi-mhz",
        type=float,
        default=_DEFAULT_SPI_MHZ,
        help=f"SPI clock in MHz (default {_DEFAULT_SPI_MHZ})",
    )
    parser.add_argument(
        "--wire-order",
        default=None,
        help="ignored (passthrough RGB); kept for compatibility",
    )


def run_bridge(
    panel_index: int,
    *,
    bind: str | None = None,
    port: int = DEFAULT_PORT,
    brightness: float = MATRIX_BRIGHTNESS,
    spi_mhz: float = _DEFAULT_SPI_MHZ,
    wire_order: str | None = None,  # unused; kept for compatibility
) -> None:
    width, height = panel_layout(panel_index)
    pixels = panel_pixel_count(panel_index)
    bind_addr = bind if bind is not None else "0.0.0.0"
    spi_hz = max(1, int(spi_mhz * 1_000_000))
    # UDP frames are already dimmed by the animator; avoid scaling twice on SPI.
    strip_cfg = panel_strip_config(
        panel_index,
        pixels,
        brightness=1.0,
        spi_max_speed_hz=spi_hz,
        passthrough=True,
        swap_rg=False,
    )

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, _SOCKET_RCVBUF)
    try:
        sock.bind((bind_addr, port))
    except OSError as exc:
        raise OSError(f"Cannot bind {bind_addr}:{port}: {exc}") from exc

    print(
        f"Panel {panel_index} bridge on {bind_addr}:{port} "
        f"({width}×{height}, {pixels} LEDs, SPI bus {strip_cfg.spi_bus}, "
        f"{spi_mhz:g} MHz, rpi5-ws2812, passthrough)"
    )
    print("Ctrl+C to stop")

    bufsize = pixels * 3 + 1
    frames: queue.Queue[bytes] = queue.Queue(maxsize=1)
    stop = threading.Event()
    receiver = threading.Thread(
        target=_receiver_loop,
        args=(sock, bufsize, pixels, panel_index, frames, stop),
        name="panel-bridge-recv",
        daemon=True,
    )
    receiver.start()

    try:
        with LedMatrix(width, height, config0=strip_cfg) as matrix:
            while True:
                rgb = frames.get()
                while True:
                    try:
                        rgb = frames.get_nowait()
                    except queue.Empty:
                        break
                matrix.show_rgb_bytes(rgb)
    finally:
        stop.set()
        sock.close()
        receiver.join(timeout=0.5)


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="UDP-to-SPI bridge (Pi panel emulates Pico firmware)"
    )
    add_bridge_args(parser)
    args = parser.parse_args(argv)
    try:
        run_bridge(
            args.panel_index,
            bind=args.bind,
            port=args.port,
            brightness=args.brightness,
            spi_mhz=args.spi_mhz,
            wire_order=args.wire_order,
        )
    except KeyboardInterrupt:
        print("\nStopped.")
        return 0
    except OSError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())