"""Picture color strip stream — extracts LED colors from captured frames."""

import threading
import time
from typing import Optional

import numpy as np

from ledgrab.core.capture.calibration import (
    AdvancedPixelMapper,
    CalibrationConfig,
    create_pixel_mapper,
)
from ledgrab.core.capture.screen_capture import extract_border_pixels
from ledgrab.storage.bindable import bfloat
from ledgrab.utils import get_logger
from ledgrab.utils.frame_limiter import FrameLimiter
from ledgrab.utils.timer import high_resolution_timer

from .base import ColorStripStream

logger = get_logger(__name__)


class PictureColorStripStream(ColorStripStream):
    """Color strip stream backed by a LiveStream (picture source).

    Runs a background thread that:
      1. Reads the latest frame from the LiveStream
      2. Extracts border pixels using the calibration's border_width
      3. Maps border pixels to LED colors via PixelMapper
      4. Applies temporal smoothing
      5. Caches the result for lock-free consumer reads

    Processing parameters can be hot-updated via update_source() without
    restarting the thread (except when the underlying LiveStream changes).
    """

    def __init__(self, live_stream, source):
        """
        Args:
            live_stream: Acquired LiveStream or Dict[str, LiveStream] for advanced mode.
                         Lifecycle managed by ColorStripStreamManager.
            source: PictureColorStripSource config
        """

        # Support both single LiveStream and dict of streams (advanced mode)
        if isinstance(live_stream, dict):
            self._live_streams = live_stream
            self._live_stream = next(iter(live_stream.values()))
        else:
            self._live_streams = {}
            self._live_stream = live_stream

        self._fps: int = 30  # internal capture rate (send FPS is on the target)
        self._frame_time: float = 1.0 / 30
        self._smoothing: float = bfloat(source.smoothing, 0.3)
        self._interpolation_mode: str = source.interpolation_mode
        self._calibration: CalibrationConfig = source.calibration
        self._pixel_mapper = create_pixel_mapper(
            self._calibration, interpolation_mode=self._interpolation_mode
        )
        cal_leds = self._calibration.get_total_leds()
        self._led_count: int = source.led_count if source.led_count > 0 else cal_leds

        # Thread-safe color cache
        self._latest_colors: Optional[np.ndarray] = None
        self._colors_lock = threading.Lock()
        self._previous_colors: Optional[np.ndarray] = None

        self._running = False
        self._thread: Optional[threading.Thread] = None
        self._last_timing: dict = {}

    @property
    def live_stream(self):
        """Public accessor for the underlying LiveStream (used by preview WebSocket)."""
        return self._live_stream

    @property
    def target_fps(self) -> int:
        return self._fps

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

    @property
    def display_index(self) -> Optional[int]:
        if self._live_streams:
            return None  # multi-source, ambiguous
        return self._live_stream.display_index

    @property
    def calibration(self) -> Optional[CalibrationConfig]:
        return self._calibration

    def start(self) -> None:
        if self._running:
            return
        self._running = True
        self._thread = threading.Thread(
            target=self._processing_loop,
            name="css-picture-stream",
            daemon=True,
        )
        self._thread.start()
        logger.info(f"PictureColorStripStream started (fps={self._fps}, leds={self._led_count})")

    def stop(self) -> None:
        self._running = False
        if self._thread:
            self._thread.join(timeout=5.0)
            if self._thread.is_alive():
                logger.warning("PictureColorStripStream thread did not terminate within 5s")
            self._thread = None
        self._latest_colors = None
        self._previous_colors = None
        logger.info("PictureColorStripStream stopped")

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

    def get_last_timing(self) -> dict:
        return dict(self._last_timing)

    def set_capture_fps(self, fps: int) -> None:
        """Update the internal capture rate. Thread-safe (read atomically by the loop)."""
        fps = max(1, min(90, fps))
        if fps != self._fps:
            self._fps = fps
            self._frame_time = 1.0 / fps
            logger.info(f"PictureColorStripStream capture FPS set to {fps}")

    def update_source(self, source) -> None:
        """Hot-update processing parameters. Thread-safe for scalar params.

        PixelMapper is rebuilt atomically if calibration or interpolation_mode changed.
        """
        from ledgrab.storage.color_strip_source import (
            AdvancedPictureColorStripSource,
            PictureColorStripSource,
        )

        if not isinstance(source, (PictureColorStripSource, AdvancedPictureColorStripSource)):
            return

        self._smoothing = bfloat(source.smoothing, 0.3)

        if (
            source.interpolation_mode != self._interpolation_mode
            or source.calibration != self._calibration
            or source.led_count != self._led_count
        ):
            self._interpolation_mode = source.interpolation_mode
            self._calibration = source.calibration
            cal_leds = source.calibration.get_total_leds()
            self._led_count = source.led_count if source.led_count > 0 else cal_leds
            self._pixel_mapper = create_pixel_mapper(
                source.calibration, interpolation_mode=source.interpolation_mode
            )
            self._previous_colors = None  # Reset smoothing history on calibration change

        logger.info("PictureColorStripStream params updated in-place")

    def _processing_loop(self) -> None:
        """Background thread: poll source, process, cache colors."""
        cached_frame = None

        # Scratch buffer pool (pre-allocated, resized when LED count changes)
        _pool_n = 0
        _frame_a = _frame_b = None  # double-buffered uint8 output
        _use_a = True
        _u16_a = _u16_b = None  # uint16 scratch for smoothing blending

        def _blend_u16(a, b, alpha_b, out):
            """Blend two uint8 arrays: out = ((256-alpha_b)*a + alpha_b*b) >> 8.

            Uses pre-allocated uint16 scratch buffers (_u16_a, _u16_b).
            """
            nonlocal _u16_a, _u16_b
            np.copyto(_u16_a, a, casting="unsafe")
            np.copyto(_u16_b, b, casting="unsafe")
            _u16_a *= 256 - alpha_b
            _u16_b *= alpha_b
            _u16_a += _u16_b
            _u16_a >>= 8
            np.copyto(out, _u16_a, casting="unsafe")

        limiter = FrameLimiter(self._fps)

        try:
            with high_resolution_timer():
                while self._running:
                    limiter.begin()
                    frame_time = self._frame_time

                    try:
                        frame = self._live_stream.get_latest_frame()

                        if frame is None or frame is cached_frame:
                            limiter.wait(frame_time)
                            continue

                        cached_frame = frame

                        t0 = time.perf_counter()

                        calibration = self._calibration
                        mapper = self._pixel_mapper

                        if isinstance(mapper, AdvancedPixelMapper):
                            # Advanced mode: gather frames from all live streams
                            frames_dict = {}
                            for ps_id, ls in self._live_streams.items():
                                f = ls.get_latest_frame()
                                if f is not None:
                                    frames_dict[ps_id] = f
                            t1 = time.perf_counter()
                            led_colors = mapper.map_lines_to_leds(frames_dict)
                        else:
                            border_pixels = extract_border_pixels(frame, calibration.border_width)
                            t1 = time.perf_counter()
                            led_colors = mapper.map_border_to_leds(border_pixels)
                        t2 = time.perf_counter()

                        # Ensure scratch pool is sized for this frame
                        target_count = self._led_count
                        _n = target_count if target_count > 0 else len(led_colors)
                        if _n > 0 and _n != _pool_n:
                            _pool_n = _n
                            _frame_a = np.empty((_n, 3), dtype=np.uint8)
                            _frame_b = np.empty((_n, 3), dtype=np.uint8)
                            _u16_a = np.empty((_n, 3), dtype=np.uint16)
                            _u16_b = np.empty((_n, 3), dtype=np.uint16)
                            self._previous_colors = None

                        # Copy/pad into double-buffered frame (avoids per-frame allocations)
                        frame_buf = _frame_a if _use_a else _frame_b
                        _use_a = not _use_a
                        n_leds = len(led_colors)
                        if _pool_n > 0:
                            if n_leds < _pool_n:
                                frame_buf[:n_leds] = led_colors
                                frame_buf[n_leds:] = 0
                            elif n_leds > _pool_n:
                                frame_buf[:] = led_colors[:_pool_n]
                            else:
                                frame_buf[:] = led_colors
                            led_colors = frame_buf

                        # Temporal smoothing (pre-allocated uint16 scratch)
                        smoothing = self.resolve("smoothing", self._smoothing)
                        if (
                            self._previous_colors is not None
                            and smoothing > 0
                            and len(self._previous_colors) == len(led_colors)
                            and _u16_a is not None
                        ):
                            _blend_u16(
                                led_colors, self._previous_colors, int(smoothing * 256), led_colors
                            )
                        t3 = time.perf_counter()

                        self._previous_colors = led_colors

                        with self._colors_lock:
                            self._latest_colors = led_colors

                        self._last_timing = {
                            "extract_ms": (t1 - t0) * 1000,
                            "map_leds_ms": (t2 - t1) * 1000,
                            "smooth_ms": (t3 - t2) * 1000,
                            "total_ms": (t3 - t0) * 1000,
                        }

                    except Exception as e:
                        logger.error(
                            f"PictureColorStripStream processing error: {e}", exc_info=True
                        )

                    limiter.wait(frame_time)
        except Exception as e:
            logger.error(f"Fatal PictureColorStripStream loop error: {e}", exc_info=True)
        finally:
            self._running = False
