"""Color strip stream management with reference counting.

PictureColorStripStreams (expensive screen capture) are shared across multiple
consumers via reference counting — processing runs once, not once per target.

Count-dependent streams (static, gradient, color cycle, effect) are NOT shared.
Each consumer gets its own instance so it can configure an independent LED count
without interfering with other targets.
"""

from dataclasses import dataclass
from typing import Dict, Optional

from wled_controller.core.processing.color_strip_stream import (
    ColorCycleColorStripStream,
    ColorStripStream,
    GradientColorStripStream,
    PictureColorStripStream,
    StaticColorStripStream,
)
from wled_controller.core.processing.processed_stream import ProcessedColorStripStream
from wled_controller.core.processing.effect_stream import EffectColorStripStream
from wled_controller.core.processing.api_input_stream import ApiInputColorStripStream
from wled_controller.core.processing.notification_stream import NotificationColorStripStream
from wled_controller.core.processing.daylight_stream import DaylightColorStripStream
from wled_controller.core.processing.candlelight_stream import CandlelightColorStripStream
from wled_controller.core.processing.game_event_stream import GameEventColorStripStream
from wled_controller.utils import get_logger

logger = get_logger(__name__)

# source_type → stream class for non-picture (non-sharable) sources
_SIMPLE_STREAM_MAP = {
    "static": StaticColorStripStream,
    "gradient": GradientColorStripStream,
    "color_cycle": ColorCycleColorStripStream,
    "effect": EffectColorStripStream,
    "api_input": ApiInputColorStripStream,
    "notification": NotificationColorStripStream,
    "daylight": DaylightColorStripStream,
    "candlelight": CandlelightColorStripStream,
    "game_event": GameEventColorStripStream,
}


@dataclass
class _ColorStripEntry:
    """Internal tracking entry for a managed color strip stream."""

    stream: ColorStripStream
    ref_count: int
    # IDs of picture sources whose LiveStreams we acquired (for release)
    picture_source_ids: list = None
    # Per-consumer target FPS values (target_id → fps)
    target_fps: Dict[str, int] = None
    # Clock ID currently acquired for this stream (for correct release)
    clock_id: Optional[str] = None
    # Value stream IDs acquired for BindableFloat properties (prop → vs_id)
    bound_vs_ids: Dict[str, str] = None

    def __post_init__(self):
        if self.picture_source_ids is None:
            self.picture_source_ids = []
        if self.target_fps is None:
            self.target_fps = {}
        if self.bound_vs_ids is None:
            self.bound_vs_ids = {}


class ColorStripStreamManager:
    """Manages ColorStripStream instances with reference counting.

    Picture streams are shared — multiple consumers reuse the same instance.
    Count-dependent streams (static, gradient, cycle, effect) are per-consumer,
    keyed by ``{css_id}:{consumer_id}``.
    """

    def __init__(
        self,
        color_strip_store,
        live_stream_manager,
        audio_capture_manager=None,
        audio_source_store=None,
        audio_template_store=None,
        sync_clock_manager=None,
        value_stream_manager=None,
        cspt_store=None,
        gradient_store=None,
        weather_manager=None,
        asset_store=None,
        game_event_bus=None,
        audio_processing_template_store=None,
    ):
        """
        Args:
            color_strip_store: ColorStripStore for resolving source configs
            live_stream_manager: LiveStreamManager for acquiring picture streams
            audio_capture_manager: AudioCaptureManager for audio-reactive sources
            audio_source_store: AudioSourceStore for resolving audio source chains
            sync_clock_manager: SyncClockManager for acquiring clock runtimes
            value_stream_manager: ValueStreamManager for per-layer brightness sources
            cspt_store: ColorStripProcessingTemplateStore for per-layer filter chains
            gradient_store: GradientStore for resolving gradient entity references
            game_event_bus: GameEventBus for game event stream subscriptions
            audio_processing_template_store: AudioProcessingTemplateStore for filter chains
        """
        self._color_strip_store = color_strip_store
        self._live_stream_manager = live_stream_manager
        self._audio_capture_manager = audio_capture_manager
        self._audio_source_store = audio_source_store
        self._audio_template_store = audio_template_store
        self._sync_clock_manager = sync_clock_manager
        self._value_stream_manager = value_stream_manager
        self._cspt_store = cspt_store
        self._gradient_store = gradient_store
        self._weather_manager = weather_manager
        self._asset_store = asset_store
        self._game_event_bus = game_event_bus
        self._audio_processing_template_store = audio_processing_template_store
        self._streams: Dict[str, _ColorStripEntry] = {}

    def _inject_clock(self, css_stream, source) -> Optional[str]:
        """Inject a SyncClockRuntime into the stream if source has clock_id.

        Returns the clock_id that was acquired, or None.
        """
        clock_id = getattr(source, "clock_id", None)
        if clock_id and self._sync_clock_manager and hasattr(css_stream, "set_clock"):
            try:
                clock_rt = self._sync_clock_manager.acquire(clock_id)
                css_stream.set_clock(clock_rt)
                logger.debug(f"Injected clock {clock_id} into stream for {source.id}")
                return clock_id
            except Exception as e:
                logger.warning(f"Could not inject clock {clock_id}: {e}")
        return None

    def _release_clock(self, source_id: str, stream, clock_id: str = None) -> None:
        """Release the clock runtime acquired for a stream.

        Args:
            source_id: CSS source ID (used as fallback to look up clock_id from store)
            stream: The stream instance (unused, kept for API compat)
            clock_id: Explicit clock_id to release. If None, looks up from store.
        """
        if not self._sync_clock_manager:
            return
        try:
            if not clock_id:
                source = self._color_strip_store.get_source(source_id)
                clock_id = getattr(source, "clock_id", None)
            if clock_id:
                self._sync_clock_manager.release(clock_id)
        except Exception as e:
            logger.debug("Sync clock release during stream cleanup: %s", e)
            pass  # source may have been deleted already

    # Properties that can be BindableFloat or BindableColor on any CSS source
    _BINDABLE_PROPS = (
        "smoothing",
        "sensitivity",
        "intensity",
        "scale",
        "speed",
        "wind_strength",
        "temperature_influence",
        "sound_volume",
        "timeout",
        "brightness",
        "duration_ms",
        "color",
        "color_peak",
        "fallback_color",
        "default_color",
    )

    def _bind_value_streams(
        self, css_stream: ColorStripStream, source, entry: _ColorStripEntry
    ) -> None:
        """Acquire ValueStreams for any bound BindableFloat/BindableColor properties."""
        if not self._value_stream_manager:
            return
        from wled_controller.storage.bindable import BindableColor, BindableFloat

        for prop in self._BINDABLE_PROPS:
            bf = getattr(source, prop, None)
            if isinstance(bf, (BindableFloat, BindableColor)) and bf.source_id:
                try:
                    vs = self._value_stream_manager.acquire(bf.source_id)
                    css_stream.set_value_stream(prop, vs)
                    entry.bound_vs_ids[prop] = bf.source_id
                    logger.debug("Bound VS %s → %s.%s", bf.source_id, source.id, prop)
                except Exception as e:
                    logger.warning(
                        "Failed to acquire VS %s for %s.%s: %s", bf.source_id, source.id, prop, e
                    )

    def _release_value_streams(self, entry: _ColorStripEntry) -> None:
        """Release all ValueStreams acquired for an entry."""
        if not self._value_stream_manager or not entry.bound_vs_ids:
            return
        for prop, vs_id in entry.bound_vs_ids.items():
            try:
                self._value_stream_manager.release(vs_id)
                entry.stream.remove_value_stream(prop)
            except Exception as e:
                logger.debug("VS release for %s: %s", vs_id, e)
        entry.bound_vs_ids.clear()

    def _resolve_key(self, css_id: str, consumer_id: str) -> str:
        """Resolve internal registry key for a (css_id, consumer_id) pair.

        Per-consumer entries (non-picture sources) take priority over shared.
        """
        if consumer_id:
            composite = f"{css_id}:{consumer_id}"
            if composite in self._streams:
                return composite
        return css_id

    def acquire(self, css_id: str, consumer_id: str = "") -> ColorStripStream:
        """Get or create a ColorStripStream for the given ColorStripSource.

        Sharable sources (picture) are shared — keyed by css_id, ref-counted.
        Non-sharable sources always get a fresh instance per consumer, keyed by
        ``css_id:consumer_id``.

        Args:
            css_id: ID of the ColorStripSource config
            consumer_id: Unique consumer identifier (target_id) — used as
                registry key for non-sharable streams.

        Returns:
            ColorStripStream instance

        Raises:
            ValueError: If ColorStripSource not found or type unsupported
            RuntimeError: If stream creation/start fails
        """
        source = self._color_strip_store.get_source(css_id)

        # Non-sharable: always create a fresh per-consumer instance
        if not source.sharable:
            if source.source_type == "audio":
                from wled_controller.core.processing.audio_stream import AudioColorStripStream

                css_stream = AudioColorStripStream(
                    source,
                    self._audio_capture_manager,
                    self._audio_source_store,
                    self._audio_template_store,
                    self._audio_processing_template_store,
                )
            elif source.source_type == "composite":
                from wled_controller.core.processing.composite_stream import (
                    CompositeColorStripStream,
                )

                css_stream = CompositeColorStripStream(
                    source, self, self._value_stream_manager, self._cspt_store
                )
            elif source.source_type == "mapped":
                from wled_controller.core.processing.mapped_stream import MappedColorStripStream

                css_stream = MappedColorStripStream(source, self)
            elif source.source_type == "processed":
                css_stream = ProcessedColorStripStream(source, self, self._cspt_store)
            elif source.source_type == "weather":
                from wled_controller.core.processing.weather_stream import WeatherColorStripStream

                css_stream = WeatherColorStripStream(source, self._weather_manager)
            else:
                stream_cls = _SIMPLE_STREAM_MAP.get(source.source_type)
                if not stream_cls:
                    raise ValueError(
                        f"Unsupported color strip source type '{source.source_type}' for {css_id}"
                    )
                css_stream = stream_cls(source)
            # Inject gradient store for palette resolution
            if self._gradient_store and hasattr(css_stream, "set_gradient_store"):
                css_stream.set_gradient_store(self._gradient_store)
            # Inject asset store for notification sound playback
            if self._asset_store and hasattr(css_stream, "set_asset_store"):
                css_stream.set_asset_store(self._asset_store)
            # Inject game event bus for game event streams
            if self._game_event_bus and hasattr(css_stream, "set_event_bus"):
                css_stream.set_event_bus(self._game_event_bus)
            # Inject sync clock runtime if source references a clock
            acquired_clock_id = self._inject_clock(css_stream, source)
            css_stream.start()
            key = f"{css_id}:{consumer_id}" if consumer_id else css_id
            entry = _ColorStripEntry(
                stream=css_stream,
                ref_count=1,
                picture_source_ids=[],
                clock_id=acquired_clock_id,
            )
            self._bind_value_streams(css_stream, source, entry)
            self._streams[key] = entry
            logger.info(f"Created {source.source_type} stream {key}")
            return css_stream

        # Sharable (picture): reuse existing stream if running
        if css_id in self._streams:
            entry = self._streams[css_id]
            entry.ref_count += 1
            logger.info(f"Reusing stream {css_id} (ref_count={entry.ref_count})")
            return entry.stream

        # Key Colors: sharable, needs a single LiveStream
        from wled_controller.storage.color_strip_source import KeyColorsColorStripSource

        if isinstance(source, KeyColorsColorStripSource):
            ps_id = source.picture_source_id
            if not ps_id:
                raise ValueError(f"Key colors source {css_id} has no picture_source_id assigned")
            try:
                live_stream = self._live_stream_manager.acquire(ps_id)
            except Exception as e:
                raise ValueError(
                    f"Failed to acquire live stream for key_colors {css_id}: {e}"
                ) from e
            try:
                from wled_controller.core.processing.kc_color_strip_stream import (
                    KeyColorsColorStripStream,
                )

                css_stream = KeyColorsColorStripStream(live_stream, source)
                css_stream.start()
            except Exception as e:
                self._live_stream_manager.release(ps_id)
                raise RuntimeError(f"Failed to start key_colors stream {css_id}: {e}") from e

            entry = _ColorStripEntry(
                stream=css_stream,
                ref_count=1,
                picture_source_ids=[ps_id],
            )
            self._bind_value_streams(css_stream, source, entry)
            self._streams[css_id] = entry
            logger.info(f"Created key_colors stream {css_id} ({len(source.rectangles)} rects)")
            return css_stream

        # Create new picture stream — needs LiveStream(s) from the capture pipeline
        from wled_controller.storage.color_strip_source import (
            PictureColorStripSource,
            AdvancedPictureColorStripSource,
        )

        if not isinstance(source, (PictureColorStripSource, AdvancedPictureColorStripSource)):
            raise ValueError(
                f"Unsupported sharable source type '{source.source_type}' for {css_id}"
            )

        # Determine required picture sources based on calibration mode
        required_ps_ids = source.calibration.get_required_picture_source_ids()
        if not required_ps_ids:
            # Simple mode: use the CSS source's single picture_source_id
            ps_id = getattr(source, "picture_source_id", "")
            if not ps_id:
                raise ValueError(f"Color strip source {css_id} has no picture_source_id assigned")
            required_ps_ids = [ps_id]

        # Acquire all required live streams (with rollback on failure)
        acquired = {}
        try:
            for ps_id in required_ps_ids:
                acquired[ps_id] = self._live_stream_manager.acquire(ps_id)
        except Exception as e:
            for ps_id in acquired:
                self._live_stream_manager.release(ps_id)
            raise ValueError(f"Failed to acquire live streams for source {css_id}: {e}") from e

        # Create stream (single LiveStream for simple, dict for advanced)
        try:
            if len(acquired) == 1 and source.calibration.mode == "simple":
                live_arg = next(iter(acquired.values()))
            else:
                live_arg = acquired
            css_stream = PictureColorStripStream(live_arg, source)
            css_stream.start()
        except Exception as e:
            for ps_id in acquired:
                self._live_stream_manager.release(ps_id)
            raise RuntimeError(
                f"Failed to start color strip stream for source {css_id}: {e}"
            ) from e

        entry = _ColorStripEntry(
            stream=css_stream,
            ref_count=1,
            picture_source_ids=list(acquired.keys()),
        )
        self._bind_value_streams(css_stream, source, entry)
        self._streams[css_id] = entry

        logger.info(f"Created picture color strip stream {css_id}")
        return css_stream

    def release(self, css_id: str, consumer_id: str = "") -> None:
        """Release a reference to a ColorStripStream.

        Decrements the reference count. When it reaches 0, stops the stream,
        releases the underlying LiveStream, and removes from registry.

        Args:
            css_id: ID of the ColorStripSource to release
            consumer_id: Consumer identifier (same as passed to acquire)
        """
        key = self._resolve_key(css_id, consumer_id)
        entry = self._streams.get(key)
        if not entry:
            logger.warning(f"Attempted to release unknown color strip stream: {key}")
            return

        entry.ref_count -= 1
        logger.debug(f"Released color strip stream {key} (ref_count={entry.ref_count})")

        if entry.ref_count <= 0:
            try:
                entry.stream.stop()
            except Exception as e:
                logger.error(f"Error stopping color strip stream {key}: {e}")

            # Release clock runtime if acquired
            if entry.clock_id:
                source_id = key.split(":")[0] if ":" in key else key
                self._release_clock(source_id, entry.stream, clock_id=entry.clock_id)

            # Release bound value streams
            self._release_value_streams(entry)

            picture_source_ids = entry.picture_source_ids
            del self._streams[key]
            logger.info(f"Removed color strip stream {key}")

            # Release all underlying live streams
            for ps_id in picture_source_ids:
                self._live_stream_manager.release(ps_id)

    def update_source(self, css_id: str, new_source) -> None:
        """Hot-update processing params on all running streams for a source.

        Applies to both the shared entry (picture) and per-consumer entries
        (count-dependent). If the picture_source_id changed, only updates
        in-place params; the underlying LiveStream is not swapped (requires
        target restart to take full effect).

        Args:
            css_id: ID of the ColorStripSource
            new_source: Updated ColorStripSource config
        """
        # Find all entries: shared key OR per-consumer keys (css_id:xxx)
        matching_keys = [k for k in self._streams if k == css_id or k.startswith(f"{css_id}:")]
        if not matching_keys:
            return  # Stream not running; config will be used on next acquire

        for key in matching_keys:
            entry = self._streams[key]
            old_clock_id = entry.clock_id
            entry.stream.update_source(new_source)
            # Hot-swap clock if clock_id changed
            if hasattr(entry.stream, "set_clock") and self._sync_clock_manager:
                new_clock_id = getattr(new_source, "clock_id", None)
                if new_clock_id:
                    if new_clock_id != old_clock_id:
                        try:
                            clock_rt = self._sync_clock_manager.acquire(new_clock_id)
                            entry.stream.set_clock(clock_rt)
                            entry.clock_id = new_clock_id
                            # Release old clock after acquiring new one
                            if old_clock_id:
                                source_id = key.split(":")[0] if ":" in key else key
                                self._release_clock(source_id, entry.stream, clock_id=old_clock_id)
                        except Exception as e:
                            logger.warning(f"Could not hot-swap clock {new_clock_id}: {e}")
                elif old_clock_id:
                    entry.stream.set_clock(None)
                    entry.clock_id = None
                    source_id = key.split(":")[0] if ":" in key else key
                    self._release_clock(source_id, entry.stream, clock_id=old_clock_id)

        # Track picture source changes for future reference counting
        from wled_controller.storage.color_strip_source import (
            PictureColorStripSource,
            AdvancedPictureColorStripSource,
        )

        if isinstance(new_source, (PictureColorStripSource, AdvancedPictureColorStripSource)):
            new_ps_ids = new_source.calibration.get_required_picture_source_ids()
            if not new_ps_ids:
                ps_id = getattr(new_source, "picture_source_id", "")
                new_ps_ids = [ps_id] if ps_id else []
            for key in matching_keys:
                entry = self._streams[key]
                if set(new_ps_ids) != set(entry.picture_source_ids):
                    logger.info(
                        f"CSS {css_id}: picture source set changed — "
                        f"restart target to use new sources"
                    )
                    break

        logger.info(f"Updated {len(matching_keys)} running stream(s) for source {css_id}")

    def refresh_audio_filter_pipelines(self, template_id: str) -> None:
        """Rebuild audio filter pipelines for any running AudioColorStripStream
        that references the given audio processing template ID.

        Called when an audio processing template is updated or deleted.
        """
        from wled_controller.core.processing.audio_stream import AudioColorStripStream

        count = 0
        for entry in self._streams.values():
            stream = entry.stream
            if isinstance(stream, AudioColorStripStream):
                if template_id in getattr(stream, "_audio_processing_template_ids", []):
                    stream._rebuild_filter_pipeline()
                    count += 1
        if count:
            logger.info(
                "Refreshed audio filter pipeline for %d stream(s) after template %s update",
                count,
                template_id,
            )

    def notify_target_fps(self, css_id: str, target_id: str, fps: int) -> None:
        """Register or update a consumer's target FPS.

        Recalculates the capture rate for PictureColorStripStreams as
        max(all consumer FPS values). Non-picture streams are unaffected.
        """
        key = self._resolve_key(css_id, target_id)
        entry = self._streams.get(key)
        if not entry:
            return
        entry.target_fps[target_id] = fps
        self._recalc_capture_fps(entry)

    def remove_target_fps(self, css_id: str, target_id: str) -> None:
        """Unregister a consumer's target FPS (e.g. on stop)."""
        key = self._resolve_key(css_id, target_id)
        entry = self._streams.get(key)
        if not entry:
            return
        entry.target_fps.pop(target_id, None)
        self._recalc_capture_fps(entry)

    def _recalc_capture_fps(self, entry: _ColorStripEntry) -> None:
        """Push max(consumer FPS) to the stream if it supports set_capture_fps."""
        if not hasattr(entry.stream, "set_capture_fps"):
            return
        if entry.target_fps:
            new_fps = max(entry.target_fps.values())
        else:
            new_fps = 30  # default when no consumers
        entry.stream.set_capture_fps(new_fps)

    def get_streams_by_source_id(self, css_id: str) -> list:
        """Return all running stream instances for a given source ID.

        Checks both the shared key (css_id) and per-consumer keys (css_id:xxx).
        Used by the REST/WS push handlers to broadcast pushed colors to all
        consumers of an api_input source.
        """
        streams = []
        for key, entry in self._streams.items():
            if key == css_id or key.startswith(f"{css_id}:"):
                streams.append(entry.stream)
        return streams

    def release_all(self) -> None:
        """Stop and remove all managed color strip streams. Called on shutdown."""
        css_ids = list(self._streams.keys())
        for css_id in css_ids:
            entry = self._streams.get(css_id)
            if entry:
                try:
                    entry.stream.stop()
                except Exception as e:
                    logger.error(f"Error stopping color strip stream {css_id}: {e}")
        self._streams.clear()
        logger.info("Released all managed color strip streams")

    def get_active_stream_ids(self) -> list:
        """Get list of active stream IDs with ref counts (for diagnostics)."""
        return [{"id": sid, "ref_count": entry.ref_count} for sid, entry in self._streams.items()]
