"""Key Colors target processor — extracts dominant colors from screen regions."""

from __future__ import annotations

import asyncio
import collections
import json
import time
from datetime import datetime, timezone
from typing import Dict, List, Optional, Tuple

import cv2
import numpy as np

from wled_controller.core.processing.live_stream import LiveStream
from wled_controller.core.capture.screen_capture import (
    calculate_average_color,
    calculate_dominant_color,
    calculate_median_color,
)
from wled_controller.core.processing.target_processor import (
    ProcessingMetrics,
    TargetContext,
    TargetProcessor,
)
from wled_controller.utils import get_logger
from wled_controller.utils.timer import high_resolution_timer

logger = get_logger(__name__)

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


# ---------------------------------------------------------------------------
# CPU-bound frame processing (runs in thread pool via asyncio.to_thread)
# ---------------------------------------------------------------------------

def _process_kc_frame(capture, rect_names, rect_bounds, calc_fn, prev_colors_arr, smoothing, brightness):
    """All CPU-bound work for one KC frame.

    Returns (colors, colors_arr, timing_ms) where:
    - colors is a dict {name: (r, g, b)}
    - colors_arr is a (N, 3) float64 array for smoothing continuity
    - timing_ms is a dict with per-stage timing in milliseconds.
    """
    t0 = time.perf_counter()

    # Downsample to working resolution — 144x fewer pixels at 1080p
    small = cv2.resize(capture.image, KC_WORK_SIZE, interpolation=cv2.INTER_AREA)

    # Extract colors for each rectangle from the small image
    n = len(rect_names)
    colors_arr = np.empty((n, 3), dtype=np.float64)
    for i, (y1, y2, x1, x2) in enumerate(rect_bounds):
        colors_arr[i] = calc_fn(small[y1:y2, x1:x2])

    t1 = time.perf_counter()

    # Vectorized smoothing on (N, 3) array
    if prev_colors_arr is not None and smoothing > 0:
        colors_arr = colors_arr * (1 - smoothing) + prev_colors_arr * smoothing

    # Apply brightness scaling
    if brightness < 1.0:
        output_arr = colors_arr * brightness
    else:
        output_arr = colors_arr

    colors_u8 = np.clip(output_arr, 0, 255).astype(np.uint8)
    t2 = time.perf_counter()

    # Build output dict
    colors = {rect_names[i]: tuple(int(c) for c in colors_u8[i]) for i in range(n)}

    timing_ms = {
        "calc_colors": (t1 - t0) * 1000,
        "smooth": (t2 - t1) * 1000,
        "total": (t2 - t0) * 1000,
    }
    return colors, colors_arr, timing_ms


# ---------------------------------------------------------------------------
# KCTargetProcessor
# ---------------------------------------------------------------------------

class KCTargetProcessor(TargetProcessor):
    """Extracts key colors from screen capture regions and broadcasts via WebSocket."""

    def __init__(
        self,
        target_id: str,
        picture_source_id: str,
        settings,  # KeyColorsSettings
        ctx: TargetContext,
    ):
        super().__init__(target_id, ctx, picture_source_id)
        self._settings = settings
        self._brightness_vs_id = settings.brightness_value_source_id if settings else ""

        # Runtime state
        self._live_stream: Optional[LiveStream] = None
        self._value_stream = None  # active brightness value stream
        self._previous_colors: Optional[Dict[str, Tuple[int, int, int]]] = None
        self._latest_colors: Optional[Dict[str, Tuple[int, int, int]]] = None
        self._ws_clients: List = []
        self._resolved_target_fps: Optional[int] = None
        self._resolved_rectangles = None

    # ----- Properties -----

    @property
    def settings(self):
        return self._settings

    # ----- Lifecycle -----

    async def start(self) -> None:
        if self._is_running:
            logger.debug(f"KC target {self._target_id} is already running")
            return

        if not self._picture_source_id:
            raise ValueError(f"KC target {self._target_id} has no picture source assigned")

        if not self._settings.pattern_template_id:
            raise ValueError(f"KC target {self._target_id} has no pattern template assigned")

        # Resolve pattern template to get rectangles
        try:
            pattern_template = self._ctx.pattern_template_store.get_template(
                self._settings.pattern_template_id
            )
        except (ValueError, AttributeError):
            raise ValueError(
                f"Pattern template {self._settings.pattern_template_id} not found"
            )

        if not pattern_template.rectangles:
            raise ValueError(
                f"Pattern template {self._settings.pattern_template_id} has no rectangles"
            )

        self._resolved_rectangles = pattern_template.rectangles

        # Acquire live stream
        try:
            live_stream = await asyncio.to_thread(
                self._ctx.live_stream_manager.acquire, self._picture_source_id
            )
            self._live_stream = live_stream
            self._resolved_target_fps = live_stream.target_fps
            logger.info(
                f"Acquired live stream for KC target {self._target_id} "
                f"(picture_source={self._picture_source_id})"
            )
        except Exception as e:
            logger.error(f"Failed to initialize live stream for KC target {self._target_id}: {e}")
            raise RuntimeError(f"Failed to initialize live stream: {e}")

        # Acquire value stream for brightness modulation (if configured)
        if self._brightness_vs_id and self._ctx.value_stream_manager:
            try:
                self._value_stream = self._ctx.value_stream_manager.acquire(
                    self._brightness_vs_id
                )
            except Exception as e:
                logger.warning(f"Failed to acquire value stream {self._brightness_vs_id}: {e}")
                self._value_stream = None

        # Reset metrics
        self._metrics = ProcessingMetrics(start_time=datetime.now(timezone.utc))
        self._previous_colors = None
        self._latest_colors = None

        # Start processing task
        self._is_running = True
        self._task = asyncio.create_task(self._processing_loop())

        logger.info(f"Started KC processing for target {self._target_id}")
        self._ctx.fire_event({"type": "state_change", "target_id": self._target_id, "processing": True})

    async def stop(self) -> None:
        if not self._is_running:
            logger.warning(f"KC processing not running for target {self._target_id}")
            return

        self._is_running = False

        # Cancel task
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                logger.debug("KC target processor task cancelled")
                pass
            self._task = None

        # Release live stream
        if self._live_stream:
            try:
                self._ctx.live_stream_manager.release(self._picture_source_id)
            except Exception as e:
                logger.warning(f"Error releasing live stream for KC target: {e}")
            self._live_stream = None

        # Release value stream
        if self._value_stream is not None and self._ctx.value_stream_manager:
            try:
                self._ctx.value_stream_manager.release(self._brightness_vs_id)
            except Exception as e:
                logger.warning(f"Error releasing value stream: {e}")
        self._value_stream = None

        logger.info(f"Stopped KC processing for target {self._target_id}")
        self._ctx.fire_event({"type": "state_change", "target_id": self._target_id, "processing": False})

    # ----- Settings -----

    def update_settings(self, settings) -> None:
        self._settings = settings
        # Keep _brightness_vs_id in sync (hot-swap handled separately)
        self._brightness_vs_id = settings.brightness_value_source_id if settings else ""
        logger.info(f"Updated KC target settings: {self._target_id}")

    def update_brightness_value_source(self, vs_id: str) -> None:
        """Hot-swap the brightness value source for a running KC target."""
        old_vs_id = self._brightness_vs_id
        self._brightness_vs_id = vs_id
        vs_mgr = self._ctx.value_stream_manager

        if not self._is_running or vs_mgr is None:
            return

        # Release old stream
        if self._value_stream is not None and old_vs_id:
            try:
                vs_mgr.release(old_vs_id)
            except Exception as e:
                logger.warning(f"Error releasing old value stream {old_vs_id}: {e}")
            self._value_stream = None

        # Acquire new stream
        if vs_id:
            try:
                self._value_stream = vs_mgr.acquire(vs_id)
            except Exception as e:
                logger.warning(f"Failed to acquire value stream {vs_id}: {e}")
                self._value_stream = None

        logger.info(f"Hot-swapped brightness VS for KC target {self._target_id}: {old_vs_id} -> {vs_id}")

    # ----- State / Metrics -----

    def get_state(self) -> dict:
        metrics = self._metrics
        return {
            "target_id": self._target_id,
            "processing": self._is_running,
            "fps_actual": round(metrics.fps_actual, 1) if self._is_running else None,
            "fps_potential": metrics.fps_potential if self._is_running else None,
            "fps_target": self._settings.fps,
            "frames_skipped": metrics.frames_skipped if self._is_running else None,
            "frames_keepalive": metrics.frames_keepalive if self._is_running else None,
            "fps_current": metrics.fps_current if self._is_running else None,
            "timing_calc_colors_ms": round(metrics.timing_calc_colors_ms, 1) if self._is_running else None,
            "timing_smooth_ms": round(metrics.timing_smooth_ms, 1) if self._is_running else None,
            "timing_broadcast_ms": round(metrics.timing_broadcast_ms, 1) if self._is_running else None,
            "timing_total_ms": round(metrics.timing_total_ms, 1) if self._is_running else None,
            "last_update": metrics.last_update,
            "errors": [metrics.last_error] if metrics.last_error else [],
            "brightness_value_source_id": self._brightness_vs_id,
        }

    def get_metrics(self) -> dict:
        metrics = self._metrics
        uptime = 0.0
        if metrics.start_time and self._is_running:
            uptime = (datetime.now(timezone.utc) - metrics.start_time).total_seconds()

        return {
            "target_id": self._target_id,
            "processing": self._is_running,
            "fps_actual": round(metrics.fps_actual, 1),
            "fps_target": self._settings.fps,
            "uptime_seconds": round(uptime, 1),
            "frames_processed": metrics.frames_processed,
            "errors_count": metrics.errors_count,
            "last_error": metrics.last_error,
            "last_update": metrics.last_update.isoformat() if metrics.last_update else None,
        }

    # ----- WebSocket -----

    def supports_websocket(self) -> bool:
        return True

    def add_ws_client(self, ws) -> None:
        self._ws_clients.append(ws)

    def remove_ws_client(self, ws) -> None:
        if ws in self._ws_clients:
            self._ws_clients.remove(ws)

    def get_latest_colors(self) -> Dict[str, Tuple[int, int, int]]:
        return self._latest_colors or {}

    # ----- Private: processing loop -----

    async def _processing_loop(self) -> None:
        """Main processing loop for key-colors extraction."""
        target_fps = self._settings.fps

        # Lookup table for interpolation mode → function (used per-frame from live settings)
        calc_fns = {
            "average": calculate_average_color,
            "median": calculate_median_color,
            "dominant": calculate_dominant_color,
        }

        frame_time = 1.0 / target_fps
        fps_samples: collections.deque = collections.deque(maxlen=10)
        timing_samples: collections.deque = collections.deque(maxlen=10)
        prev_frame_time_stamp = time.perf_counter()
        prev_capture = None
        last_broadcast_time = 0.0
        send_timestamps: collections.deque = collections.deque()

        rectangles = self._resolved_rectangles

        # Pre-compute pixel bounds at working resolution (160x90)
        kc_w, kc_h = KC_WORK_SIZE
        rect_names = [r.name for r in rectangles]
        rect_bounds = []
        for rect in rectangles:
            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)
            rect_bounds.append((px_y, px_y + px_h, px_x, px_x + px_w))
        prev_colors_arr = None

        logger.info(
            f"KC processing loop started for target {self._target_id} "
            f"(fps={target_fps}, rects={len(rectangles)})"
        )

        try:
            with high_resolution_timer():
                while self._is_running:
                    loop_start = time.perf_counter()

                    try:
                        capture = self._live_stream.get_latest_frame()

                        if capture is None:
                            await asyncio.sleep(frame_time)
                            continue

                        # Skip processing if the frame hasn't changed
                        if capture is prev_capture:
                            # Keepalive: re-broadcast last colors
                            if self._latest_colors and (loop_start - last_broadcast_time) >= 1.0:
                                await self._broadcast_colors(self._latest_colors)
                                last_broadcast_time = time.perf_counter()
                                send_timestamps.append(last_broadcast_time)
                                self._metrics.frames_keepalive += 1
                            self._metrics.frames_skipped += 1
                            now_ts = time.perf_counter()
                            while send_timestamps and send_timestamps[0] < now_ts - 1.0:
                                send_timestamps.popleft()
                            self._metrics.fps_current = len(send_timestamps)
                            await asyncio.sleep(frame_time)
                            continue
                        prev_capture = capture

                        # Read settings fresh each frame so hot updates (brightness,
                        # smoothing, interpolation_mode) take effect without restart.
                        s = self._settings
                        calc_fn = calc_fns.get(s.interpolation_mode, calculate_average_color)

                        # Effective brightness: static setting * value stream
                        eff_brightness = s.brightness
                        vs = self._value_stream
                        if vs is not None:
                            eff_brightness *= vs.get_value()

                        # CPU-bound work in thread pool
                        colors, colors_arr, frame_timing = await asyncio.to_thread(
                            _process_kc_frame,
                            capture, rect_names, rect_bounds, calc_fn,
                            prev_colors_arr, s.smoothing, eff_brightness,
                        )

                        prev_colors_arr = colors_arr
                        self._latest_colors = dict(colors)

                        # Broadcast to WebSocket clients
                        t_broadcast_start = time.perf_counter()
                        await self._broadcast_colors(colors)
                        broadcast_ms = (time.perf_counter() - t_broadcast_start) * 1000
                        last_broadcast_time = time.perf_counter()
                        send_timestamps.append(last_broadcast_time)

                        # Per-stage timing (rolling average over last 10 frames)
                        frame_timing["broadcast"] = broadcast_ms
                        timing_samples.append(frame_timing)
                        n = len(timing_samples)
                        self._metrics.timing_calc_colors_ms = sum(s["calc_colors"] for s in timing_samples) / n
                        self._metrics.timing_smooth_ms = sum(s["smooth"] for s in timing_samples) / n
                        self._metrics.timing_broadcast_ms = sum(s["broadcast"] for s in timing_samples) / n
                        self._metrics.timing_total_ms = sum(s["total"] for s in timing_samples) / n + broadcast_ms

                        # Update metrics
                        self._metrics.frames_processed += 1
                        self._metrics.last_update = datetime.now(timezone.utc)

                        # Calculate actual FPS
                        now = time.perf_counter()
                        interval = now - prev_frame_time_stamp
                        prev_frame_time_stamp = now
                        fps_samples.append(1.0 / interval if interval > 0 else 0)
                        self._metrics.fps_actual = sum(fps_samples) / len(fps_samples)

                        # Potential FPS
                        processing_time = now - loop_start
                        self._metrics.fps_potential = 1.0 / processing_time if processing_time > 0 else 0

                        # fps_current
                        while send_timestamps and send_timestamps[0] < now - 1.0:
                            send_timestamps.popleft()
                        self._metrics.fps_current = len(send_timestamps)

                    except Exception as e:
                        self._metrics.errors_count += 1
                        self._metrics.last_error = str(e)
                        logger.error(f"KC processing error for {self._target_id}: {e}", exc_info=True)

                    # Throttle to target FPS
                    elapsed = time.perf_counter() - loop_start
                    remaining = frame_time - elapsed
                    if remaining > 0:
                        await asyncio.sleep(remaining)

        except asyncio.CancelledError:
            logger.info(f"KC processing loop cancelled for target {self._target_id}")
            raise
        except Exception as e:
            logger.error(f"Fatal error in KC processing loop for target {self._target_id}: {e}")
            self._is_running = False
            self._ctx.fire_event({"type": "state_change", "target_id": self._target_id, "processing": False, "crashed": True})
            raise
        finally:
            logger.info(f"KC processing loop ended for target {self._target_id}")

    async def _broadcast_colors(self, colors: Dict[str, Tuple[int, int, int]]) -> None:
        """Broadcast extracted colors to WebSocket clients (concurrent sends)."""
        if not self._ws_clients:
            return

        message = json.dumps({
            "type": "colors_update",
            "target_id": self._target_id,
            "colors": {
                name: {"r": c[0], "g": c[1], "b": c[2]}
                for name, c in colors.items()
            },
            "timestamp": datetime.now(timezone.utc).isoformat(),
        })

        async def _send_safe(ws):
            try:
                await ws.send_text(message)
                return True
            except Exception as e:
                logger.debug("KC WS send failed: %s", e)
                return False

        clients = list(self._ws_clients)
        results = await asyncio.gather(*[_send_safe(ws) for ws in clients])

        for ws, ok in zip(clients, results):
            if not ok and ws in self._ws_clients:
                self._ws_clients.remove(ws)
