#!/usr/bin/env python3
"""Positioning and sync test for multiple Pico panels.

Runs a sequence of patterns on all panels (or one via --panel-index):
  identify   - panel index digit on each board (check physical mapping)
  corners    - red/green/blue/yellow corners (check orientation)
  crosshair  - center row + column
  columns    - sweep the same column index on every panel together
  rows       - sweep rows together
  flash      - simultaneous white flashes (check sync)
  pulse      - in-sync brightness pulse on all panels
  width      - sweep columns near right edge (find 38 vs 39 vs 40)
  all        - full sequence (default)
"""

from __future__ import annotations

import argparse
import math
import socket
import sys
import time

from leds.panel_udp import (
    HeadlessMatrix,
    PanelTarget,
    add_panel_args,
    format_panel_targets,
    panel_targets_from_args,
    send_frame,
)
from leds.text import draw_text_centered

PANEL_COLORS = (
    (255, 80, 80),
    (80, 255, 80),
    (80, 160, 255),
    (255, 200, 80),
    (200, 80, 255),
)

CORNER_COLORS = {
    "top-left": (255, 0, 0),
    "top-right": (0, 255, 0),
    "bottom-left": (0, 0, 255),
    "bottom-right": (255, 255, 0),
}


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


def _push(
    sock: socket.socket,
    targets: list[PanelTarget],
    matrices: list[tuple[PanelTarget, HeadlessMatrix]],
    port: int,
    panel_id: int | None,
) -> None:
    """Send all panel frames in a tight burst for best sync."""
    for target, matrix in matrices:
        send_frame(sock, target.host, port, matrix.rgb_bytes(), panel_id)


def _clear_all(
    sock: socket.socket,
    targets: list[PanelTarget],
    port: int,
    brightness: float,
    panel_id: int | None,
) -> None:
    matrices = _matrices(targets, brightness)
    for _, matrix in matrices:
        matrix.clear()
    _push(sock, targets, matrices, port, panel_id)


def test_identify(
    sock: socket.socket,
    targets: list[PanelTarget],
    port: int,
    brightness: float,
    panel_id: int | None,
    pause: float,
) -> None:
    print("Test: identify — each panel shows its index (0–4)")
    matrices = _matrices(targets, brightness)
    for target, matrix in matrices:
        matrix.clear()
        color = PANEL_COLORS[target.index % len(PANEL_COLORS)]
        draw_text_centered(matrix, str(target.index), color)
        print(f"  panel {target.index} -> {target.host} ({target.width}×{target.height})")
    _push(sock, targets, matrices, port, panel_id)
    time.sleep(pause)


def test_corners(
    sock: socket.socket,
    targets: list[PanelTarget],
    port: int,
    brightness: float,
    panel_id: int | None,
    pause: float,
) -> None:
    print("Test: corners — TL=red TR=green BL=blue BR=yellow (all panels)")
    matrices = _matrices(targets, brightness)
    for target, matrix in matrices:
        w, h = target.width, target.height
        matrix.clear()
        points = {
            "top-left": (0, 0),
            "top-right": (w - 1, 0),
            "bottom-left": (0, h - 1),
            "bottom-right": (w - 1, h - 1),
        }
        for name, (x, y) in points.items():
            matrix[x, y] = CORNER_COLORS[name]
    _push(sock, targets, matrices, port, panel_id)
    time.sleep(pause)


def test_crosshair(
    sock: socket.socket,
    targets: list[PanelTarget],
    port: int,
    brightness: float,
    panel_id: int | None,
    pause: float,
) -> None:
    print("Test: crosshair — center row + column (white)")
    matrices = _matrices(targets, brightness)
    for target, matrix in matrices:
        w, h = target.width, target.height
        matrix.clear()
        mid_y = h // 2
        mid_x = w // 2
        for x in range(w):
            matrix[x, mid_y] = (255, 255, 255)
        for y in range(h):
            matrix[mid_x, y] = (200, 200, 200)
    _push(sock, targets, matrices, port, panel_id)
    time.sleep(pause)


def test_columns(
    sock: socket.socket,
    targets: list[PanelTarget],
    port: int,
    brightness: float,
    panel_id: int | None,
    pause: float,
) -> None:
    max_width = max(t.width for t in targets)
    print(f"Test: column sweep — x=0..{max_width - 1} on all panels together")
    for x in range(max_width):
        matrices = _matrices(targets, brightness)
        for target, matrix in matrices:
            matrix.clear()
            if x >= target.width:
                continue
            hue = int(255 * x / max(max_width - 1, 1))
            for y in range(target.height):
                matrix[x, y] = (hue, 80, 255 - hue)
        _push(sock, targets, matrices, port, panel_id)
        time.sleep(pause)
    _clear_all(sock, targets, port, brightness, panel_id)


def test_rows(
    sock: socket.socket,
    targets: list[PanelTarget],
    port: int,
    brightness: float,
    panel_id: int | None,
    pause: float,
) -> None:
    height = targets[0].height if targets else 9
    print(f"Test: row sweep — y=0..{height - 1} on all panels together")
    for y in range(height):
        matrices = _matrices(targets, brightness)
        for target, matrix in matrices:
            matrix.clear()
            for x in range(target.width):
                matrix[x, y] = (255, 255, 255)
        _push(sock, targets, matrices, port, panel_id)
        time.sleep(pause)
    _clear_all(sock, targets, port, brightness, panel_id)


def test_flash(
    sock: socket.socket,
    targets: list[PanelTarget],
    port: int,
    brightness: float,
    panel_id: int | None,
    pause: float,
    count: int,
) -> None:
    print(f"Test: sync flash — {count} white flashes on all panels")
    for n in range(count):
        matrices = _matrices(targets, brightness)
        for _, matrix in matrices:
            matrix.fill((255, 255, 255))
        _push(sock, targets, matrices, port, panel_id)
        time.sleep(pause)
        _clear_all(sock, targets, port, brightness, panel_id)
        time.sleep(pause / 2)
        print(f"  flash {n + 1}/{count}")


def test_pulse(
    sock: socket.socket,
    targets: list[PanelTarget],
    port: int,
    brightness: float,
    panel_id: int | None,
    seconds: float,
    fps: float,
) -> None:
    print(f"Test: sync pulse — {seconds:.0f}s solid red fade (all panels)")
    frames = max(int(seconds * fps), 1)
    for frame in range(frames):
        t = frame / frames
        level = 0.15 + 0.85 * (0.5 - 0.5 * math.cos(t * 6.28318))
        color = (int(255 * level), 0, 0)
        matrices = _matrices(targets, min(brightness * level, 1.0))
        for _, matrix in matrices:
            matrix.fill(color)
        _push(sock, targets, matrices, port, panel_id)
        time.sleep(1.0 / fps)
    _clear_all(sock, targets, port, brightness, panel_id)


def test_chain_ends(
    sock: socket.socket,
    targets: list[PanelTarget],
    port: int,
    brightness: float,
    panel_id: int | None,
    pause: float,
) -> None:
    print("Test: chain ends — LED 0=red, last LED=green (per panel)")
    matrices = _matrices(targets, brightness)
    for target, matrix in matrices:
        matrix.clear()
        matrix[0, 0] = (255, 0, 0)
        last = target.pixels - 1
        x = last % target.width
        y = last // target.width
        matrix[x, y] = (0, 255, 0)
        print(f"  panel {target.index}: index 0 and {last} ({x},{y})")
    _push(sock, targets, matrices, port, panel_id)
    time.sleep(pause)


def test_width(
    sock: socket.socket,
    targets: list[PanelTarget],
    port: int,
    brightness: float,
    panel_id: int | None,
    pause: float,
    col_start: int,
    col_end: int,
) -> None:
    probe_w = col_end + 1
    probe_leds = probe_w * targets[0].height if targets else probe_w * 9
    print("Test: width — one full column at a time (right-edge probe)")
    print(f"  Probing columns {col_start}–{col_end} (sending {probe_w}×9 = {probe_leds} LEDs)")
    print("  Flash firmware with at least NUM_LEDS={} first, e.g.:".format(probe_leds))
    print(f"    make deploy PANEL_ID=<n> NUM_LEDS={probe_leds}")
    print("  The last column that lights a real LED → width = column + 1")
    print("  (column 37 lit → width 38; column 38 lit → width 39)\n")
    for x in range(col_start, col_end + 1):
        burst: list[tuple[PanelTarget, HeadlessMatrix]] = []
        for target in targets:
            matrix = HeadlessMatrix(probe_w, target.height, brightness=brightness)
            matrix.clear()
            for y in range(target.height):
                matrix[x, y] = (255, 255, 255)
            probe_target = PanelTarget(target.index, target.host, probe_w, target.height)
            burst.append((probe_target, matrix))
        for probe_target, matrix in burst:
            send_frame(sock, probe_target.host, port, matrix.rgb_bytes(), panel_id)
        print(f"  column {x}  → width {x + 1} if this is the last real column")
        time.sleep(pause)
    _clear_all(sock, targets, port, brightness, panel_id)


TESTS = {
    "identify": test_identify,
    "corners": test_corners,
    "crosshair": test_crosshair,
    "columns": test_columns,
    "rows": test_rows,
    "flash": test_flash,
    "pulse": test_pulse,
    "ends": test_chain_ends,
    "width": test_width,
}

ALL_ORDER = ("identify", "corners", "ends", "crosshair", "columns", "rows", "flash", "pulse")


def main() -> int:
    parser = argparse.ArgumentParser(description="Multi-panel positioning and sync test")
    add_panel_args(parser)
    parser.add_argument(
        "--test",
        choices=[*TESTS.keys(), "all"],
        default="all",
        help="which test to run (default: full sequence)",
    )
    parser.add_argument("--pause", type=float, default=1.2, help="hold time per static pattern")
    parser.add_argument("--column-pause", type=float, default=0.12, help="delay per column in sweep")
    parser.add_argument("--row-pause", type=float, default=0.25, help="delay per row in sweep")
    parser.add_argument("--flashes", type=int, default=3, help="count for sync flash test")
    parser.add_argument("--pulse-seconds", type=float, default=4.0)
    parser.add_argument("--fps", type=float, default=30.0)
    parser.add_argument("--col-start", type=int, default=32, help="width test: first column")
    parser.add_argument("--col-end", type=int, default=42, help="width test: last column")
    args = parser.parse_args()

    targets = panel_targets_from_args(args)
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    print(f"Panel sync test -> {format_panel_targets(targets, args.port)}")
    print("Ctrl+C to stop\n")

    tests = ALL_ORDER if args.test == "all" else (args.test,)

    try:
        for name in tests:
            fn = TESTS[name]
            if name == "flash":
                fn(sock, targets, args.port, args.brightness, args.panel_id, args.pause, args.flashes)
            elif name == "pulse":
                fn(sock, targets, args.port, args.brightness, args.panel_id, args.pulse_seconds, args.fps)
            elif name == "columns":
                fn(sock, targets, args.port, args.brightness, args.panel_id, args.column_pause)
            elif name == "rows":
                fn(sock, targets, args.port, args.brightness, args.panel_id, args.row_pause)
            elif name == "width":
                fn(sock, targets, args.port, args.brightness, args.panel_id, args.pause, args.col_start, args.col_end)
            else:
                fn(sock, targets, args.port, args.brightness, args.panel_id, args.pause)
            print()
        print("Done.")
    except KeyboardInterrupt:
        print("\nStopped.")
    except OSError as exc:
        print(f"Network error: {exc}", file=sys.stderr)
        return 1
    finally:
        _clear_all(sock, targets, args.port, args.brightness, args.panel_id)
        sock.close()
    return 0


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