#!/usr/bin/env python3
"""Bench-test a WS2812B ring: ramp brightness, then set levels interactively.

Wiring (Pi Zero 2W):
  ring DI  -> GPIO13 (physical pin 33)   [PWM1 -- matches ws2812b.py]
  ring GND -> any Pi GND (pin 34 is adjacent) + PSU GND if externally powered
  ring 5V  -> external 5V supply for full brightness; the Pi header 5V rail
              is only safe for low-brightness testing (see --max-brightness)

rpi_ws281x needs root for DMA/PWM access:
  sudo python3 led_ring_test.py
  sudo python3 led_ring_test.py --count 32 --pin 13 --color 255,255,255
"""

from __future__ import annotations

import argparse
import sys
import time


def parse_color(text: str):
    parts = [int(p) for p in text.split(",")]
    if len(parts) != 3 or not all(0 <= p <= 255 for p in parts):
        raise argparse.ArgumentTypeError("color must be R,G,B with each 0-255")
    return tuple(parts)


def main() -> int:
    parser = argparse.ArgumentParser(description="WS2812B ring brightness test")
    parser.add_argument("--pin", type=int, default=13, help="BCM GPIO for data (default 13 = phys pin 33)")
    parser.add_argument("--count", type=int, default=32, help="number of LEDs (default 32)")
    parser.add_argument("--color", type=parse_color, default=(255, 255, 255), help="R,G,B (default white)")
    parser.add_argument(
        "--max-brightness", type=int, default=64,
        help="cap for the ramp/interactive values, 0-255 (default 64 -- ~0.5A "
             "white on 32 LEDs, safe from the Pi header; raise only on an "
             "external supply)",
    )
    parser.add_argument("--no-ramp", action="store_true", help="skip the ramp, go straight to interactive")
    parser.add_argument("--rgbw", action="store_true",
                        help="drive as SK6812 RGBW (4 channels/LED) -- try this if only "
                             "part of the ring lights in scrambled colors")
    args = parser.parse_args()

    from rpi_ws281x import PixelStrip, Color
    import _rpi_ws281x as ws

    # PWM channel is fixed by the pin: GPIO12/18 -> 0, GPIO13/19 -> 1.
    channel = 1 if args.pin in (13, 19) else 0
    strip_type = ws.SK6812_STRIP_RGBW if args.rgbw else ws.WS2811_STRIP_GRB
    strip = PixelStrip(args.count, args.pin, freq_hz=800000, dma=10, invert=False,
                       brightness=args.max_brightness, channel=channel,
                       strip_type=strip_type)
    strip.begin()

    color = Color(*args.color)

    def show(brightness: int) -> None:
        strip.setBrightness(brightness)
        for i in range(strip.numPixels()):
            strip.setPixelColor(i, color)
        strip.show()

    try:
        if not args.no_ramp:
            print(f"Ramping 0 -> {args.max_brightness} -> 0 on GPIO{args.pin}, {args.count} LEDs...")
            steps = list(range(0, args.max_brightness + 1, 4))
            for b in steps + steps[::-1]:
                show(b)
                time.sleep(0.05)

        print(f"Enter brightness 0-{args.max_brightness} (q to quit):")
        while True:
            line = input("> ").strip().lower()
            if line in ("q", "quit", "exit", ""):
                break
            try:
                b = int(line)
            except ValueError:
                print("not a number")
                continue
            if not 0 <= b <= args.max_brightness:
                print(f"out of range 0-{args.max_brightness} (raise with --max-brightness "
                      "if you're on an external supply)")
                continue
            show(b)
            amps = args.count * 0.060 * (b / 255)
            print(f"brightness {b} (~{amps:.2f}A at full white)")
    except (KeyboardInterrupt, EOFError):
        pass
    finally:
        show(0)
        print("\nLEDs off.")

    return 0


if __name__ == "__main__":
    sys.exit(main())