"""Tests for the shared edge-to-LED interpolation kernels."""

from __future__ import annotations

import numpy as np

from ledgrab.core.capture.edge_interpolation import (
    average_edge_to_leds,
    fallback_edge_to_leds,
)

# ---------------------------------------------------------------------------
# average_edge_to_leds
# ---------------------------------------------------------------------------


def test_average_top_edge_returns_uint8_per_led():
    """A solid-colour top edge maps to a solid-colour LED array."""
    edge = np.full((4, 100, 3), [200, 100, 50], dtype=np.uint8)
    cache: dict = {}
    out = average_edge_to_leds(edge, "top", led_count=10, cache=cache, cache_key="top")
    assert out.shape == (10, 3)
    assert out.dtype == np.uint8
    # Each LED is the exact mean of its segment — which is uniform 200/100/50.
    assert np.all(out[:, 0] == 200)
    assert np.all(out[:, 1] == 100)
    assert np.all(out[:, 2] == 50)


def test_average_left_edge_uses_axis_1():
    """A solid-colour LEFT edge (collapsing across columns) produces the same."""
    edge = np.full((100, 4, 3), [10, 20, 30], dtype=np.uint8)
    cache: dict = {}
    out = average_edge_to_leds(edge, "left", led_count=5, cache=cache, cache_key="left")
    assert out.shape == (5, 3)
    assert np.all(out[:, 0] == 10)
    assert np.all(out[:, 1] == 20)
    assert np.all(out[:, 2] == 30)


def test_average_cache_is_reused_across_calls():
    """Second call with the same (edge_len, led_count) does not rebuild the cache."""
    edge = np.full((4, 100, 3), 128, dtype=np.uint8)
    cache: dict = {}
    average_edge_to_leds(edge, "top", led_count=10, cache=cache, cache_key="top")
    entry_first = cache["top"]
    average_edge_to_leds(edge, "top", led_count=10, cache=cache, cache_key="top")
    entry_second = cache["top"]
    # Same tuple object → cache reused
    assert entry_first is entry_second


def test_average_cache_rebuilds_when_signature_changes():
    edge = np.full((4, 100, 3), 128, dtype=np.uint8)
    cache: dict = {}
    average_edge_to_leds(edge, "top", led_count=10, cache=cache, cache_key="top")
    entry_first = cache["top"]
    average_edge_to_leds(edge, "top", led_count=20, cache=cache, cache_key="top")
    entry_second = cache["top"]
    assert entry_first is not entry_second
    assert entry_second[1] == 20


def test_average_caches_keyed_independently():
    """Two cache keys produce independent cache entries."""
    edge = np.full((4, 100, 3), 128, dtype=np.uint8)
    cache: dict = {}
    average_edge_to_leds(edge, "top", led_count=10, cache=cache, cache_key="a")
    average_edge_to_leds(edge, "top", led_count=10, cache=cache, cache_key="b")
    assert "a" in cache and "b" in cache
    assert cache["a"] is not cache["b"]


def test_average_gradient_edge_produces_gradient_output():
    """A 1D ramp across the edge averages to a ramp across LEDs."""
    edge = np.zeros((1, 100, 3), dtype=np.uint8)
    edge[0, :, 0] = np.arange(100)  # red ramps 0..99 across width
    cache: dict = {}
    out = average_edge_to_leds(edge, "top", led_count=10, cache=cache, cache_key="top")
    # First LED: average of positions 0..9 → 4.5 → uint8 → 4
    assert int(out[0, 0]) == 4
    # Last LED: average of positions 90..99 → 94.5 → uint8 → 94
    assert int(out[-1, 0]) == 94
    # Monotonic increase
    assert all(out[i, 0] <= out[i + 1, 0] for i in range(9))


# ---------------------------------------------------------------------------
# fallback_edge_to_leds
# ---------------------------------------------------------------------------


def test_fallback_uses_provided_calc_color():
    """A trivial calc_color that returns a constant maps every LED to that constant."""
    edge = np.full((4, 100, 3), [200, 100, 50], dtype=np.uint8)
    out = fallback_edge_to_leds(
        edge, "top", led_count=10, calc_color=lambda seg: np.array([7, 8, 9])
    )
    assert out.shape == (10, 3)
    assert out.dtype == np.uint8
    assert np.all(out == [7, 8, 9])


def test_fallback_segments_track_edge_position():
    """Each LED's calc_color receives the segment slice for its position."""
    edge = np.zeros((1, 10, 3), dtype=np.uint8)
    edge[0, :, 0] = np.arange(10)

    seen_segments = []

    def _spy(seg):
        seen_segments.append(seg[:, :, 0].copy().tolist())
        return seg[0, 0]  # return the first pixel

    fallback_edge_to_leds(edge, "top", led_count=5, calc_color=_spy)
    # Five segments, two pixels each
    assert len(seen_segments) == 5
    assert seen_segments[0] == [[0, 1]]
    assert seen_segments[-1] == [[8, 9]]


def test_fallback_left_edge_segments_track_height():
    edge = np.zeros((10, 1, 3), dtype=np.uint8)
    edge[:, 0, 0] = np.arange(10)
    fallback_edge_to_leds(edge, "left", led_count=2, calc_color=lambda s: s[0, 0])
    # Should not raise; covered above shape-wise.
