"""Shared helpers for color strip streams (gradient computation, etc.)."""

import numpy as np


def _compute_gradient_colors(stops: list, led_count: int, easing: str = "linear") -> np.ndarray:
    """Compute an (led_count, 3) uint8 array from gradient color stops.

    Each stop: {"position": float 0–1, "color": [R,G,B], "color_right": [R,G,B] | absent}

    Interpolation:
      Sort stops by position.  For each LED at relative position p = i/(N-1):
        p ≤ first stop  → first stop primary color
        p ≥ last stop   → last stop right color (if bidirectional) else primary
        else find surrounding stops A (≤p) and B (>p):
          left_color  = A["color_right"] if present, else A["color"]
          right_color = B["color"]
          t = (p - A.pos) / (B.pos - A.pos)
          color = lerp(left_color, right_color, eased(t))
    """
    if led_count <= 0:
        led_count = 1

    if not stops:
        return np.zeros((led_count, 3), dtype=np.uint8)

    sorted_stops = sorted(stops, key=lambda s: float(s.get("position", 0)))

    def _color(stop: dict, side: str = "left") -> np.ndarray:
        if side == "right":
            cr = stop.get("color_right")
            if cr and isinstance(cr, list) and len(cr) == 3:
                return np.array(cr, dtype=np.float32)
        c = stop.get("color", [255, 255, 255])
        return np.array(
            c if isinstance(c, list) and len(c) == 3 else [255, 255, 255], dtype=np.float32
        )

    # Vectorized: compute all LED positions at once
    positions = np.linspace(0, 1, led_count) if led_count > 1 else np.array([0.0])
    result = np.zeros((led_count, 3), dtype=np.float32)

    # Extract stop positions and colors into arrays
    n_stops = len(sorted_stops)
    stop_positions = np.array([float(s.get("position", 0)) for s in sorted_stops], dtype=np.float32)

    # Pre-compute left/right colors for each stop
    left_colors = np.array([_color(s, "left") for s in sorted_stops], dtype=np.float32)
    right_colors = np.array([_color(s, "right") for s in sorted_stops], dtype=np.float32)

    # LEDs before first stop
    mask_before = positions <= stop_positions[0]
    result[mask_before] = left_colors[0]

    # LEDs after last stop
    mask_after = positions >= stop_positions[-1]
    result[mask_after] = right_colors[-1]

    # LEDs between stops — vectorized per segment
    mask_between = ~mask_before & ~mask_after
    if np.any(mask_between):
        between_pos = positions[mask_between]
        # np.searchsorted finds the right stop index for each LED
        idx = np.searchsorted(stop_positions, between_pos, side="right") - 1
        idx = np.clip(idx, 0, n_stops - 2)

        a_pos = stop_positions[idx]
        b_pos = stop_positions[idx + 1]
        span = b_pos - a_pos
        t = np.where(span > 0, (between_pos - a_pos) / span, 0.0)

        # Apply easing to interpolation parameter
        if easing == "ease_in_out":
            t = t * t * (3.0 - 2.0 * t)
        elif easing == "cubic":
            t = np.where(t < 0.5, 4.0 * t * t * t, 1.0 - (-2.0 * t + 2.0) ** 3 / 2.0)
        elif easing == "step":
            steps = float(max(2, n_stops))
            t = np.round(t * steps) / steps

        a_colors = right_colors[idx]  # A's right color
        b_colors = left_colors[idx + 1]  # B's left color
        result[mask_between] = a_colors + t[:, np.newaxis] * (b_colors - a_colors)

    return np.clip(result, 0, 255).astype(np.uint8)
