"""
On-device test that exercises the Chase pattern via Presets.

Usage (from pico/ dir or project root with adjusted paths):

    mpremote connect <device> cp src/*.py :
    mpremote connect <device> cp src/patterns/*.py :patterns
    mpremote connect <device> cp lib/*.py :
    mpremote connect <device> cp test/test_chase_via_presets.py :
    mpremote connect <device> run test_chase_via_presets.py
"""

import utime

from presets import Presets, Preset


def snapshot_strip_colors(presets, strip_idx=0, max_leds=32):
    """Return a list of (r,g,b) tuples for the first max_leds of the given strip."""
    strip = presets.strips[strip_idx]
    num = min(strip.num_leds, max_leds)
    out = []
    for i in range(num):
        o = i * 3
        g = strip.ar[o]
        r = strip.ar[o + 1]
        b = strip.ar[o + 2]
        out.append((r, g, b))
    return out


def expected_chase_color(i, num_leds, step_count, color0, color1, n1, n2, n3, n4):
    """Mirror the position logic from patterns/chase.py for a single logical LED."""
    segment_length = n1 + n2

    if step_count % 2 == 0:
        position = (step_count // 2) * (n3 + n4) + n3
    else:
        position = ((step_count + 1) // 2) * (n3 + n4)

    max_pos = num_leds + segment_length
    position = position % max_pos
    if position < 0:
        position += max_pos

    relative_pos = (i - position) % segment_length
    if relative_pos < 0:
        relative_pos = (relative_pos + segment_length) % segment_length

    return color0 if relative_pos < n1 else color1


def test_chase_single_step_via_presets():
    presets = Presets()

    num_leds = presets.num_leds
    if num_leds <= 0:
        print("No strips; skipping chase test.")
        return

    # Simple alternating colors with known lengths.
    base_color0 = (10, 0, 0)
    base_color1 = (0, 0, 20)

    # Use full brightness so apply_brightness is identity.
    brightness = 255

    n1 = 2
    n2 = 3
    # Same step size on even/odd for easier reasoning.
    n3 = 1
    n4 = 1

    data = {
        "p": "chase",
        "c": [base_color0, base_color1],
        "b": brightness,
        "d": 0,
        "a": False,  # single-step mode
        "n1": n1,
        "n2": n2,
        "n3": n3,
        "n4": n4,
    }

    name = "test_chase_pattern"
    preset = Preset(data)
    presets.presets[name] = preset

    # Select and run one tick; this should render exactly one chase frame for step 0.
    presets.select(name, step=0)
    presets.tick()

    # Colors after brightness scaling (driver.apply_brightness is used in the pattern).
    color0 = presets.apply_brightness(base_color0, brightness)
    color1 = presets.apply_brightness(base_color1, brightness)

    # Snapshot first few LEDs of strip 0 and compare against expected pattern for step 0.
    colors = snapshot_strip_colors(presets, strip_idx=0, max_leds=16)
    step_count = 0

    for i, actual in enumerate(colors):
        expected = expected_chase_color(
            i, num_leds, step_count, color0, color1, n1, n2, n3, n4
        )
        assert (
            actual == expected
        ), "LED %d: got %r, expected %r" % (i, actual, expected)

    print("test_chase_single_step_via_presets: OK")


def test_chase_multiple_steps_via_presets():
    """Render several steps and verify pattern advances correctly."""
    presets = Presets()

    num_leds = presets.num_leds
    if num_leds <= 0:
        print("No strips; skipping chase multi-step test.")
        return

    base_color0 = (10, 0, 0)
    base_color1 = (0, 0, 20)
    brightness = 255

    n1 = 2
    n2 = 3
    n3 = 1
    n4 = 1

    data = {
        "p": "chase",
        "c": [base_color0, base_color1],
        "b": brightness,
        "d": 0,
        "a": False,
        "n1": n1,
        "n2": n2,
        "n3": n3,
        "n4": n4,
    }

    name = "test_chase_pattern_multi"
    preset = Preset(data)
    presets.presets[name] = preset

    color0 = presets.apply_brightness(base_color0, brightness)
    color1 = presets.apply_brightness(base_color1, brightness)

    # In non-auto mode (a=False), the Chase pattern advances one step per
    # invocation of the generator, and Presets is expected to call select()
    # again for each beat. Emulate that here by re-selecting with an
    # explicit step value for each frame we want to test.
    for step_count in range(4):
        presets.select(name, step=step_count)
        presets.tick()
        colors = snapshot_strip_colors(presets, strip_idx=0, max_leds=16)

        for i, actual in enumerate(colors):
            expected = expected_chase_color(
                i, num_leds, step_count, color0, color1, n1, n2, n3, n4
            )
            assert (
                actual == expected
            ), "step %d, LED %d: got %r, expected %r" % (
                step_count,
                i,
                actual,
                expected,
            )

    print("test_chase_multiple_steps_via_presets: OK")


def main():
    test_chase_single_step_via_presets()
    test_chase_multiple_steps_via_presets()
    # Give a brief pause so message is visible if run interactively.
    utime.sleep_ms(100)


if __name__ == "__main__":
    main()