#!/usr/bin/env python3
"""Drive-test all 8 WS2812 strips on led-driver-8 over UDP.

Protocol (port 50007):
  - 1 + N×3 bytes: strip_id (0–7, or 255=all) + RGB
  - 8 bytes PIN\\0 + strip + gpio + len_lo + len_hi

Defaults match firmware board_config.h strip map.

Usage:
  python3 examples/test_strips.py
  python3 examples/test_strips.py --host 10.1.1.10 --leds 64
  python3 examples/test_strips.py --host 10.1.1.10 --all-only
"""

from __future__ import annotations

import argparse
import socket
import time

DEFAULT_STRIPS = (
    # (strip_id, gpio)
    (0, 18),
    (1, 19),
    (2, 20),
    (3, 21),
    (4, 22),
    (5, 26),
    (6, 27),
    (7, 28),
)

IDENTITY_COLOURS = (
    (255, 0, 0),  # red
    (0, 255, 0),  # green
    (0, 0, 255),  # blue
    (255, 255, 0),  # yellow
    (255, 0, 255),  # magenta
    (0, 255, 255),  # cyan
    (255, 128, 0),  # orange
    (255, 255, 255),  # white
)

ALL_COLOURS = (
    ("RED", (255, 0, 0)),
    ("GREEN", (0, 255, 0)),
    ("BLUE", (0, 0, 255)),
    ("WHITE", (255, 255, 255)),
)


def pin_bind(sock: socket.socket, addr: tuple[str, int], strip: int, gpio: int, leds: int) -> None:
    payload = bytes(
        [
            ord("P"),
            ord("I"),
            ord("N"),
            0,
            strip & 0xFF,
            gpio & 0xFF,
            leds & 0xFF,
            (leds >> 8) & 0xFF,
        ]
    )
    sock.sendto(payload, addr)


def send_solid(
    sock: socket.socket,
    addr: tuple[str, int],
    strip_id: int,
    leds: int,
    rgb: tuple[int, int, int],
) -> None:
    sock.sendto(bytes([strip_id & 0xFF]) + bytes(rgb) * leds, addr)


def off_all(sock: socket.socket, addr: tuple[str, int], leds: int) -> None:
    for strip, _gpio in DEFAULT_STRIPS:
        send_solid(sock, addr, strip, leds, (0, 0, 0))


def main() -> int:
    parser = argparse.ArgumentParser(description="Test all 8 led-driver-8 strips over UDP")
    parser.add_argument("--host", default="10.1.1.10", help="Panel IP (default: 10.1.1.10)")
    parser.add_argument("--port", type=int, default=50007, help="UDP RGB port")
    parser.add_argument("--leds", type=int, default=64, help="LEDs lit per strip (capacity test)")
    parser.add_argument("--pause", type=float, default=0.7, help="Seconds between steps")
    parser.add_argument(
        "--all-only",
        action="store_true",
        help="Skip per-strip identity; only flash all strips together",
    )
    parser.add_argument("--no-bind", action="store_true", help="Skip PIN bind commands")
    args = parser.parse_args()

    if args.leds < 1:
        parser.error("--leds must be >= 1")

    addr = (args.host, args.port)
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

    print(f"Testing {len(DEFAULT_STRIPS)} strips @ {args.host}:{args.port} ({args.leds} LEDs each)")

    try:
        if not args.no_bind:
            for strip, gpio in DEFAULT_STRIPS:
                pin_bind(sock, addr, strip, gpio, args.leds)
                print(f"  bind strip {strip} GP{gpio}")
                time.sleep(0.02)

        if not args.all_only:
            print("  per-strip identity colours")
            for (strip, gpio), rgb in zip(DEFAULT_STRIPS, IDENTITY_COLOURS):
                off_all(sock, addr, args.leds)
                send_solid(sock, addr, strip, args.leds, rgb)
                print(f"    strip {strip} GP{gpio} -> {rgb}")
                time.sleep(args.pause)

        print("  all strips together (strip_id=255)")
        for label, rgb in ALL_COLOURS:
            print(f"    {label}")
            send_solid(sock, addr, 255, args.leds, rgb)
            time.sleep(args.pause)

        off_all(sock, addr, args.leds)
        print("Done.")
    except OSError as exc:
        print(f"Network error: {exc}")
        return 1
    finally:
        sock.close()
    return 0


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