#!/usr/bin/env python3
"""Step-by-step drive test for Pico panel firmware over UDP."""

from __future__ import annotations

import argparse
import sys
import time

from leds.array_config import MATRIX_HEIGHT, MATRIX_PIXELS, MATRIX_WIDTH, matrix_pixel_index
from leds.panel_udp import PanelClient, add_panel_args, format_panel_targets, panel_targets_from_args


def test_flash(client: PanelClient, pause: float) -> None:
    print("Test: RGB flash (red, green, blue, off)")
    for name, color in (("red", (255, 0, 0)), ("green", (0, 255, 0)), ("blue", (0, 0, 255))):
        print(f"  {name}")
        client.fill(color)
        time.sleep(pause)
    client.clear()
    time.sleep(pause / 2)


def test_corners(client: PanelClient, pause: float) -> None:
    target = client.targets[0]
    width, height = target.width, target.height
    print("Test: matrix corners (left → right rows)")
    corners = [
        ("top-left", matrix_pixel_index(0, 0, width, height), (255, 0, 0)),
        ("top-right", matrix_pixel_index(width - 1, 0, width, height), (0, 255, 0)),
        ("bottom-left", matrix_pixel_index(0, height - 1, width, height), (0, 0, 255)),
        ("bottom-right", matrix_pixel_index(width - 1, height - 1, width, height), (255, 255, 0)),
    ]
    for label, index, color in corners:
        if index >= client.pixel_count:
            print(f"  skip {label} (index {index})")
            continue
        print(f"  {label} -> LED {index}")
        client.set_indices([index], color)
        time.sleep(pause)
    client.clear()
    time.sleep(pause / 2)


def test_rows(client: PanelClient, pause: float) -> None:
    target = client.targets[0]
    width, height = target.width, target.height
    print("Test: one row at a time (white)")
    width = min(width, client.pixel_count)
    rows = min(height, (client.pixel_count + width - 1) // width)
    for y in range(rows):
        indices = [matrix_pixel_index(x, y, width, height) for x in range(width)]
        indices = [i for i in indices if i < client.pixel_count]
        print(f"  row {y} ({len(indices)} LEDs)")
        client.set_indices(indices, (255, 255, 255))
        time.sleep(pause)
    client.clear()
    time.sleep(pause / 2)


def test_chase(client: PanelClient, pause: float, steps: int | None) -> None:
    count = steps if steps is not None else min(client.pixel_count, 60)
    print(f"Test: chase ({count} steps, cyan)")
    for i in range(count):
        client.set_indices([i % client.pixel_count], (0, 255, 255))
        time.sleep(pause)
    client.clear()
    time.sleep(pause / 2)


def test_first_last(client: PanelClient, pause: float) -> None:
    print("Test: first and last LED in chain")
    last = client.pixel_count - 1
    print(f"  LED 0 = red, LED {last} = green")
    pixels = [(0, 0, 0)] * client.pixel_count
    pixels[0] = (255, 0, 0)
    pixels[last] = (0, 255, 0)
    client.send_pixels(pixels)
    time.sleep(pause)
    client.clear()
    time.sleep(pause / 2)


TESTS = {
    "flash": lambda c, p, _s: test_flash(c, p),
    "first-last": lambda c, p, _s: test_first_last(c, p),
    "corners": lambda c, p, _s: test_corners(c, p),
    "rows": lambda c, p, _s: test_rows(c, p),
    "chase": test_chase,
}


def run_all(client: PanelClient, pause: float, chase_steps: int | None) -> None:
    test_flash(client, pause)
    test_first_last(client, pause)
    if client.pixel_count >= client.targets[0].width:
        test_corners(client, pause)
        test_rows(client, max(pause * 0.6, 0.15))
    test_chase(client, max(pause * 0.25, 0.05), chase_steps)
    client.clear()
    print("Done.")


def main() -> int:
    parser = argparse.ArgumentParser(description="Drive-test Pico panel LEDs over UDP")
    add_panel_args(parser)
    parser.add_argument("--pause", type=float, default=0.8)
    parser.add_argument("--test", choices=[*TESTS.keys(), "all"], default="all")
    parser.add_argument("--chase-steps", type=int, default=None)
    parser.add_argument("--color", default=None, metavar="R,G,B")
    args = parser.parse_args()

    if args.count is not None and args.count < 1:
        print("Error: --count must be at least 1", file=sys.stderr)
        return 1

    targets = panel_targets_from_args(args)
    print(f"Panel drive test -> {format_panel_targets(targets, args.port)}")

    try:
        with PanelClient.from_args(args) as client:
            if args.color:
                parts = [int(p.strip()) for p in args.color.split(",")]
                if len(parts) != 3:
                    raise SystemExit("--color must be R,G,B")
                client.fill((parts[0], parts[1], parts[2]))
                print(f"Sent solid color {tuple(parts)}")
                return 0

            if args.test == "all":
                run_all(client, args.pause, args.chase_steps)
            else:
                TESTS[args.test](client, args.pause, args.chase_steps)
                client.clear()
                print("Done.")
    except KeyboardInterrupt:
        print("\nStopped.")
    except OSError as exc:
        print(f"Network error: {exc}", file=sys.stderr)
        return 1
    return 0


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