portal/leds/panel_udp.py

"""UDP helpers for Pico panel firmware."""

from __future__ import annotations

import argparse
import os
import socket
import time
from dataclasses import dataclass
from typing import Iterable, Sequence, Tuple

from leds.animations import ANIMATIONS, DEFAULT_FPS
from leds.array_config import (
    MATRIX_HEIGHT,
    MATRIX_PIXELS,
    MATRIX_WIDTH,
    PANEL_COUNT,
    matrix_pixel_index,
    panel_layout,
    panel_pixel_count,
)
from leds.colors import dim

PANEL_IP_NET = "10.1.1"
PANEL_IP_BASE = 10
PANEL_MAX_FPS = 25.0

RGB = Tuple[int, int, int]

DEFAULT_PORT = int(os.environ.get("PORTAL_PANEL_PORT", "50007"))


def panel_ip(panel_id: int) -> str:
    return f"{PANEL_IP_NET}.{PANEL_IP_BASE + panel_id}"


DEFAULT_HOST = panel_ip(0)
DEFAULT_PANEL_HOSTS = [panel_ip(i) for i in range(PANEL_COUNT)]


@dataclass(frozen=True)
class PanelTarget:
    index: int
    host: str
    width: int
    height: int

    @property
    def pixels(self) -> int:
        return self.width * self.height

    def label(self) -> str:
        if self.index >= 0:
            return f"{self.host} ({self.width}×{self.height}, panel {self.index})"
        return f"{self.host} ({self.width}×{self.height})"


def panel_hosts(panel_index: int | None = None, count: int = PANEL_COUNT) -> list[str]:
    return [t.host for t in resolve_panel_targets(panel_index=panel_index, panels=count)]


def resolve_panel_targets(
    host: str | None = None,
    panel_index: int | None = None,
    panels: int = PANEL_COUNT,
    *,
    width: int | None = None,
    height: int | None = None,
) -> list[PanelTarget]:
    if panel_index is not None:
        w, h = panel_layout(panel_index)
        if width is not None:
            w = width
        if height is not None:
            h = height
        return [PanelTarget(panel_index, panel_ip(panel_index), w, h)]
    if host:
        w = width if width is not None else MATRIX_WIDTH
        h = height if height is not None else MATRIX_HEIGHT
        return [PanelTarget(-1, host, w, h)]
    return [
        PanelTarget(i, panel_ip(i), *panel_layout(i))
        for i in range(panels)
    ]


def format_panel_targets(targets: Sequence[PanelTarget], port: int) -> str:
    if len(targets) == 1:
        return f"{targets[0].label()}:{port}"
    layouts = ", ".join(f"{t.width}×{t.height}" for t in targets)
    return f"{len(targets)} panels ({targets[0].host} … {targets[-1].host}):{port} [{layouts}]"


def add_panel_network_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument(
        "--host",
        default=None,
        help="Single panel IP override (not needed; default uses 10.1.1.10–14)",
    )
    parser.add_argument(
        "--panel-index",
        type=int,
        default=None,
        choices=range(PANEL_COUNT),
        metavar="N",
        help=f"Drive only panel N (0–{PANEL_COUNT - 1}); default all {PANEL_COUNT}",
    )
    parser.add_argument(
        "--panels",
        type=int,
        default=PANEL_COUNT,
        help=f"Number of panels when driving all (default {PANEL_COUNT})",
    )
    parser.add_argument("--port", type=int, default=DEFAULT_PORT)
    parser.add_argument("--panel-id", type=int, default=None, help="Prefix UDP frames with panel id")


def add_panel_args(parser: argparse.ArgumentParser) -> None:
    add_panel_network_args(parser)
    parser.add_argument(
        "--width",
        type=int,
        default=None,
        help="Override panel width (default from leds/array_config.py per panel)",
    )
    parser.add_argument(
        "--height",
        type=int,
        default=None,
        help=f"Override panel height (default {MATRIX_HEIGHT})",
    )
    parser.add_argument(
        "--count",
        type=int,
        default=None,
        help="LEDs per frame (default: panel width × height)",
    )
    parser.add_argument("--brightness", type=float, default=0.35)


def panel_targets_from_args(args: argparse.Namespace) -> list[PanelTarget]:
    return resolve_panel_targets(
        args.host,
        args.panel_index,
        args.panels,
        width=getattr(args, "width", None),
        height=getattr(args, "height", None),
    )


def panel_count_from_args(args: argparse.Namespace) -> int:
    if args.count is not None:
        return args.count
    if args.panel_index is not None:
        return panel_pixel_count(args.panel_index)
    if args.width is not None:
        h = args.height if args.height is not None else MATRIX_HEIGHT
        return args.width * h
    return MATRIX_PIXELS


class HeadlessMatrix:
    """Matrix buffer for UDP senders — no local SPI/PIO."""

    def __init__(
        self,
        width: int = MATRIX_WIDTH,
        height: int = MATRIX_HEIGHT,
        *,
        brightness: float = 0.25,
    ) -> None:
        self.width = width
        self.height = height
        self.brightness = brightness
        self._pixels: list[RGB] = [(0, 0, 0)] * (width * height)

    def _index_at(self, x: int, y: int) -> int:
        return matrix_pixel_index(x, y, self.width, self.height)

    def __len__(self) -> int:
        return len(self._pixels)

    def __getitem__(self, pos: Tuple[int, int]) -> RGB:
        return self._pixels[self._index_at(*pos)]

    def __setitem__(self, pos: Tuple[int, int], color: RGB) -> None:
        self._pixels[self._index_at(*pos)] = color

    def fill(self, color: RGB) -> None:
        self._pixels = [color] * len(self._pixels)

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

    def set_index(self, index: int, color: RGB) -> None:
        if 0 <= index < len(self._pixels):
            self._pixels[index] = color

    def rgb_bytes(self) -> bytes:
        out = bytearray(len(self._pixels) * 3)
        for i, (r, g, b) in enumerate(self._pixels):
            r, g, b = dim((r, g, b), self.brightness)
            base = i * 3
            out[base], out[base + 1], out[base + 2] = r, g, b
        return bytes(out)


def solid_frame_bytes(
    color: RGB,
    count: int = MATRIX_PIXELS,
    brightness: float = 1.0,
) -> bytes:
    r, g, b = dim(color, brightness)
    return bytes((r, g, b)) * count


def pixels_to_frame_bytes(
    pixels: Iterable[RGB],
    brightness: float = 1.0,
) -> bytes:
    out = bytearray()
    for r, g, b in pixels:
        r, g, b = dim((r, g, b), brightness)
        out.extend((r, g, b))
    return bytes(out)


def send_frame(
    sock: socket.socket,
    host: str,
    port: int,
    rgb: bytes,
    panel_id: int | None,
) -> None:
    payload = (bytes([panel_id & 0xFF]) + rgb) if panel_id is not None else rgb
    sock.sendto(payload, (host, port))


def send_solid_to_targets(
    sock: socket.socket,
    targets: Sequence[PanelTarget],
    port: int,
    color: RGB,
    brightness: float,
    panel_id: int | None,
) -> None:
    for target in targets:
        send_frame(
            sock,
            target.host,
            port,
            solid_frame_bytes(color, target.pixels, brightness),
            panel_id,
        )


class PanelClient:
    """Send frames to one or more portal panel firmware targets."""

    def __init__(
        self,
        targets: PanelTarget | Sequence[PanelTarget] | str,
        port: int = DEFAULT_PORT,
        *,
        pixel_count: int | None = None,
        panel_id: int | None = None,
        brightness: float = 0.35,
    ) -> None:
        if isinstance(targets, PanelTarget):
            self.targets = [targets]
        elif isinstance(targets, str):
            count = pixel_count or MATRIX_PIXELS
            w = MATRIX_WIDTH
            h = count // w if count % w == 0 else MATRIX_HEIGHT
            self.targets = [PanelTarget(-1, targets, w, h)]
        else:
            self.targets = list(targets)
        self.port = port
        self.panel_id = panel_id
        self.brightness = brightness
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

    @classmethod
    def from_args(cls, args: argparse.Namespace) -> PanelClient:
        return cls(
            panel_targets_from_args(args),
            args.port,
            pixel_count=panel_count_from_args(args),
            panel_id=args.panel_id,
            brightness=args.brightness,
        )

    @property
    def pixel_count(self) -> int:
        if len(self.targets) != 1:
            raise ValueError("pixel_count requires a single panel target")
        return self.targets[0].pixels

    @property
    def hosts(self) -> list[str]:
        return [t.host for t in self.targets]

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

    def __enter__(self) -> PanelClient:
        return self

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

    def send_bytes(self, rgb: bytes) -> None:
        if len(self.targets) != 1:
            raise ValueError("send_bytes requires a single panel target")
        target = self.targets[0]
        if len(rgb) != target.pixels * 3:
            raise ValueError(f"expected {target.pixels * 3} bytes, got {len(rgb)}")
        send_frame(self._sock, target.host, self.port, rgb, self.panel_id)

    def fill(self, color: RGB = (0, 0, 0)) -> None:
        for target in self.targets:
            send_frame(
                self._sock,
                target.host,
                self.port,
                solid_frame_bytes(color, target.pixels, self.brightness),
                self.panel_id,
            )

    def send_pixels(self, pixels: list[RGB]) -> None:
        if len(self.targets) != 1:
            raise ValueError("send_pixels requires a single panel target")
        if len(pixels) != self.pixel_count:
            raise ValueError(f"expected {self.pixel_count} pixels, got {len(pixels)}")
        self.send_bytes(pixels_to_frame_bytes(pixels, self.brightness))

    def set_indices(self, indices: Iterable[int], color: RGB) -> None:
        pixels = [(0, 0, 0)] * self.pixel_count
        for i in indices:
            if 0 <= i < self.pixel_count:
                pixels[i] = color
        self.send_pixels(pixels)

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


def _panel_matrices(
    targets: Sequence[PanelTarget],
    brightness: float,
) -> list[tuple[PanelTarget, HeadlessMatrix]]:
    return [(t, HeadlessMatrix(t.width, t.height, brightness=brightness)) for t in targets]


def run_panel_playlist(
    targets: Sequence[PanelTarget],
    port: int,
    playlist: list[str],
    *,
    duration: float = 12.0,
    fps_override: float | None = None,
    brightness: float = 0.25,
    panel_id: int | None = None,
    loop: bool = True,
) -> None:
    matrices = _panel_matrices(targets, brightness)
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        while True:
            for name in playlist:
                draw = ANIMATIONS[name]
                fps = min(fps_override or DEFAULT_FPS.get(name, 30), PANEL_MAX_FPS)
                delay = 1.0 / fps
                frames = max(int(duration * fps), 1)
                print(f"  {name} ({fps:.0f} fps, {frames} frames)")
                for frame in range(frames):
                    t0 = time.monotonic()
                    for target, matrix in matrices:
                        draw(matrix, frame)
                        send_frame(sock, target.host, port, matrix.rgb_bytes(), panel_id)
                    elapsed = time.monotonic() - t0
                    if elapsed < delay:
                        time.sleep(delay - elapsed)
            if not loop:
                break
            print("  --- loop ---")
        for target, matrix in matrices:
            matrix.clear()
            send_frame(sock, target.host, port, matrix.rgb_bytes(), panel_id)
    finally:
        sock.close()


def run_panel_animation_loop(
    targets: Sequence[PanelTarget],
    port: int,
    animation: str,
    *,
    fps: float | None = None,
    brightness: float = 0.25,
    panel_id: int | None = None,
) -> None:
    draw = ANIMATIONS[animation]
    rate = min(fps or DEFAULT_FPS.get(animation, 30), PANEL_MAX_FPS)
    interval = 1.0 / rate
    matrices = _panel_matrices(targets, brightness)
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    frame = 0
    try:
        while True:
            t0 = time.monotonic()
            for target, matrix in matrices:
                draw(matrix, frame)
                send_frame(sock, target.host, port, matrix.rgb_bytes(), panel_id)
            frame += 1
            elapsed = time.monotonic() - t0
            if elapsed < interval:
                time.sleep(interval - elapsed)
    finally:
        for target, matrix in matrices:
            matrix.clear()
            send_frame(sock, target.host, port, matrix.rgb_bytes(), panel_id)
        sock.close()


def animation_playlist(animation: str) -> list[str]:
    names = list(ANIMATIONS.keys())
    if animation == "all":
        # Skip solid — 1 fps looks frozen on UDP panels; start with motion.
        return [n for n in names if n != "solid"]
    return [animation]