"""Key Colors color strip stream — extracts dominant colors from screen rectangles.

Produces an np.ndarray(N, 3) where N = number of rectangles.
Each "LED" is the average/median/dominant color of one screen region.
"""

from __future__ import annotations

import threading
import time
from typing import TYPE_CHECKING, List, Optional

import cv2
import numpy as np

from wled_controller.core.capture.screen_capture import (
    calculate_average_color,
    calculate_dominant_color,
    calculate_median_color,
)
from wled_controller.core.processing.color_strip_stream import ColorStripStream
from wled_controller.storage.bindable import bfloat
from wled_controller.utils import get_logger
from wled_controller.utils.timer import high_resolution_timer

if TYPE_CHECKING:
    from wled_controller.core.processing.live_stream import LiveStream
    from wled_controller.storage.color_strip_source import KeyColorsColorStripSource

logger = get_logger(__name__)

KC_WORK_SIZE = (160, 90)  # (width, height) — small enough for fast color calc

_CALC_FNS = {
    "average": calculate_average_color,
    "median": calculate_median_color,
    "dominant": calculate_dominant_color,
}


class KeyColorsColorStripStream(ColorStripStream):
    """Streams N colors extracted from screen rectangles."""

    def __init__(
        self,
        live_stream: LiveStream,
        source: KeyColorsColorStripSource,
    ) -> None:
        self._live_stream = live_stream
        self._source = source

        # Pre-compute rectangle pixel bounds at KC_WORK_SIZE
        kc_w, kc_h = KC_WORK_SIZE
        self._rect_names: List[str] = []
        self._rect_bounds: List[tuple] = []
        for rect in source.rectangles:
            self._rect_names.append(rect.name)
            px_x = max(0, int(rect.x * kc_w))
            px_y = max(0, int(rect.y * kc_h))
            px_w = max(1, int(rect.width * kc_w))
            px_h = max(1, int(rect.height * kc_h))
            px_x = min(px_x, kc_w - 1)
            px_y = min(px_y, kc_h - 1)
            px_w = min(px_w, kc_w - px_x)
            px_h = min(px_h, kc_h - px_y)
            self._rect_bounds.append((px_y, px_y + px_h, px_x, px_x + px_w))

        n = len(source.rectangles)
        self._led_count = n
        self._latest_colors: Optional[np.ndarray] = None
        self._colors_lock = threading.Lock()
        self._running = False
        self._thread: Optional[threading.Thread] = None

    # ── Public interface (matches ColorStripStream) ──

    @property
    def led_count(self) -> int:
        return self._led_count

    @property
    def target_fps(self) -> int:
        return self._live_stream.target_fps if self._live_stream else 10

    @property
    def is_animated(self) -> bool:
        return True

    def get_latest_colors(self) -> Optional[np.ndarray]:
        with self._colors_lock:
            return self._latest_colors

    def start(self) -> None:
        if self._running:
            return
        self._running = True
        self._thread = threading.Thread(
            target=self._processing_loop, daemon=True, name=f"kc-css-{self._source.id[:8]}"
        )
        self._thread.start()

    def stop(self) -> None:
        self._running = False
        if self._thread:
            self._thread.join(timeout=2.0)
            self._thread = None
        self._latest_colors = None

    def update_source(self, source: KeyColorsColorStripSource) -> None:
        """Hot-update source config (rectangles, interpolation, smoothing, brightness)."""
        self._source = source
        # Recompute rectangle bounds
        kc_w, kc_h = KC_WORK_SIZE
        self._rect_names = []
        self._rect_bounds = []
        for rect in source.rectangles:
            self._rect_names.append(rect.name)
            px_x = max(0, int(rect.x * kc_w))
            px_y = max(0, int(rect.y * kc_h))
            px_w = max(1, int(rect.width * kc_w))
            px_h = max(1, int(rect.height * kc_h))
            px_x = min(px_x, kc_w - 1)
            px_y = min(px_y, kc_h - 1)
            px_w = min(px_w, kc_w - px_x)
            px_h = min(px_h, kc_h - px_y)
            self._rect_bounds.append((px_y, px_y + px_h, px_x, px_x + px_w))
        self._led_count = len(source.rectangles)

    # ── Private: processing loop ──

    def _processing_loop(self) -> None:
        """Background thread: capture → extract → smooth → cache."""
        prev_capture = None
        prev_colors_arr: Optional[np.ndarray] = None
        frame_time = 1.0 / max(1, self.target_fps)

        logger.info(f"KC CSS stream started: {self._source.id} ({len(self._rect_names)} rects)")

        with high_resolution_timer():
            while self._running:
                try:
                    capture = self._live_stream.get_latest_frame()
                    if capture is None or capture is prev_capture:
                        time.sleep(frame_time)
                        continue
                    prev_capture = capture

                    # Read source config (hot-update safe)
                    src = self._source
                    calc_fn = _CALC_FNS.get(src.interpolation_mode, calculate_average_color)

                    # Downsample
                    small = cv2.resize(capture.image, KC_WORK_SIZE, interpolation=cv2.INTER_AREA)

                    # Extract colors per rectangle
                    n = len(self._rect_names)
                    if n == 0:
                        time.sleep(frame_time)
                        continue

                    colors_arr = np.empty((n, 3), dtype=np.float64)
                    for i, (y1, y2, x1, x2) in enumerate(self._rect_bounds):
                        colors_arr[i] = calc_fn(small[y1:y2, x1:x2])

                    # Temporal smoothing
                    smoothing = self.resolve("smoothing", bfloat(src.smoothing, 0.3))
                    if (
                        prev_colors_arr is not None
                        and smoothing > 0
                        and prev_colors_arr.shape == colors_arr.shape
                    ):
                        colors_arr = colors_arr * (1 - smoothing) + prev_colors_arr * smoothing
                    prev_colors_arr = colors_arr

                    # Apply brightness
                    brightness = self.resolve("brightness", bfloat(src.brightness, 1.0))
                    if brightness < 1.0:
                        output = colors_arr * brightness
                    else:
                        output = colors_arr

                    result = np.clip(output, 0, 255).astype(np.uint8)

                    with self._colors_lock:
                        self._latest_colors = result

                    time.sleep(frame_time)

                except Exception as e:
                    logger.error(f"KC CSS stream error: {e}")
                    time.sleep(0.5)

        logger.info(f"KC CSS stream stopped: {self._source.id}")
