"""Stateful decoding of encoded camera packets into raw image frames.

The framework keeps the concrete FFmpeg dependency optional.  Importing this
module therefore never imports PyAV; the default backend does so lazily when
the first encoded key frame is received.  Tests and alternative deployments
can inject another backend through ``backend_factory``.
"""

from __future__ import annotations

import threading
from collections.abc import Callable, Iterable, Mapping
from concurrent.futures import Executor
from dataclasses import dataclass
from importlib import import_module
from typing import Any, Protocol

from cmvr_edge_ai.contracts import ImageFrame
from cmvr_edge_ai.core.component import (
    ComponentContext,
    Emission,
    HealthReport,
    Operator,
)
from cmvr_edge_ai.core.envelope import Envelope
from cmvr_edge_ai.workers.thread import run_blocking


class VideoDecodeError(RuntimeError):
    """A recoverable malformed/discontinuous encoded-video packet."""


@dataclass(frozen=True, slots=True)
class DecodedVideoFrame:
    """One packed image produced by a decoder backend."""

    data: bytes
    width: int
    height: int
    pixel_format: str = "BGR8"

    def __post_init__(self) -> None:
        if not isinstance(self.data, bytes):
            raise TypeError("decoded video frame data must be bytes")
        if self.width <= 0 or self.height <= 0:
            raise ValueError("decoded video frame dimensions must be positive")
        if not isinstance(self.pixel_format, str) or not self.pixel_format.strip():
            raise ValueError("decoded video frame pixel_format must not be empty")


class VideoDecoderBackend(Protocol):
    """Synchronous stateful decoder owned by one operator instance."""

    def decode(self, data: bytes) -> Iterable[DecodedVideoFrame]:
        """Consume one encoded packet and return zero or more complete images."""

    def close(self) -> None:
        """Release native decoder resources; implementations must be idempotent."""


VideoDecoderBackendFactory = Callable[[str], VideoDecoderBackend]


class _PyAvDecoderBackend:
    """Default Annex-B H.264/H.265 decoder backed by optional PyAV."""

    def __init__(self, codec: str) -> None:
        try:
            av = import_module("av")
        except ImportError as exc:
            raise RuntimeError(
                "PyAV is required to decode encoded video; install the optional "
                "video dependencies (for example: pip install -e '.[video]')"
            ) from exc

        self._av = av
        self._codec: Any = av.CodecContext.create(codec, "r")

    def decode(self, data: bytes) -> tuple[DecodedVideoFrame, ...]:
        if self._codec is None:
            raise RuntimeError("video decoder backend is closed")
        try:
            frames = self._codec.decode(self._av.Packet(data))
        except Exception as exc:
            ffmpeg_error = getattr(
                getattr(self._av, "error", None), "FFmpegError", None
            )
            if isinstance(ffmpeg_error, type) and isinstance(exc, ffmpeg_error):
                raise VideoDecodeError(str(exc)) from exc
            raise

        decoded: list[DecodedVideoFrame] = []
        for frame in frames:
            image = frame.to_ndarray(format="bgr24")
            decoded.append(
                DecodedVideoFrame(
                    data=image.tobytes(order="C"),
                    width=int(frame.width),
                    height=int(frame.height),
                    pixel_format="BGR8",
                )
            )
        return tuple(decoded)

    def close(self) -> None:
        # PyAV CodecContext does not expose a stable explicit close API across
        # supported versions.  Dropping the final reference releases it.
        self._codec = None


class VideoDecoderOperator(Operator):
    """Decode one encoded stream while recovering at key-frame boundaries.

    A decoder context is valid only for one source/session/codec/dimension
    identity and a continuous sequence.  Any identity change or sequence gap
    closes that context and suppresses packets until the next key frame.  Raw
    frames are emitted unchanged and never import the optional decoder library.
    """

    def __init__(
        self,
        node_id: str,
        params: Mapping[str, Any],
        *,
        backend_factory: VideoDecoderBackendFactory | None = None,
    ) -> None:
        if not isinstance(node_id, str) or not node_id.strip():
            raise ValueError("video decoder node_id must not be empty")
        if params:
            unknown = ", ".join(sorted(str(key) for key in params))
            raise ValueError(f"video decoder does not support parameters: {unknown}")

        self._node_id = node_id
        self._backend_factory = backend_factory or _PyAvDecoderBackend
        self._backend: VideoDecoderBackend | None = None
        self._backend_codec: str | None = None
        self._backend_lock = threading.Lock()
        self._executor: Executor | None = None
        self._stream_identity: tuple[str, str | None, str, int, int] | None = None
        self._last_sequence: int | None = None
        self._waiting_for_keyframe = True
        self._decoded_frames = 0
        self._dropped_until_keyframe = 0
        self._resets = 0
        self._decode_errors = 0

    async def setup(self, context: ComponentContext) -> None:
        executor = context.metadata.get("thread_executor")
        if executor is not None and not isinstance(executor, Executor):
            raise TypeError("component context thread_executor must be an Executor")
        self._executor = executor

    async def process(
        self,
        envelope: Envelope[Any],
        input_port: str = "input",
    ) -> Emission | tuple[Emission, ...] | None:
        del input_port
        payload = envelope.payload
        if not isinstance(payload, ImageFrame):
            raise TypeError(
                f"{self._node_id} expected ImageFrame, got " f"{type(payload).__name__}"
            )

        codec = _normalize_codec(payload.codec)
        if codec is None:
            if self._stream_identity is not None or self._backend is not None:
                await self._reset_stream()
            return Emission("frames", envelope)

        if payload.width <= 0 or payload.height <= 0:
            raise ValueError("encoded image dimensions must be positive")
        if not payload.data:
            raise ValueError("encoded image data must not be empty")

        identity = (
            envelope.source_id,
            envelope.session_id,
            codec,
            payload.width,
            payload.height,
        )
        identity_changed = self._stream_identity != identity
        sequence_gap = (
            not identity_changed
            and self._last_sequence is not None
            and envelope.sequence != self._last_sequence + 1
        )
        if identity_changed or sequence_gap:
            await self._reset_stream(identity=identity)

        self._stream_identity = identity
        self._last_sequence = envelope.sequence
        if self._waiting_for_keyframe and not payload.is_key_frame:
            self._dropped_until_keyframe += 1
            return None

        try:
            frames = await run_blocking(
                self._decode_sync,
                codec,
                payload.data,
                executor=self._executor,
            )
        except VideoDecodeError:
            self._decode_errors += 1
            await self._reset_stream(identity=identity)
            self._last_sequence = envelope.sequence
            return None

        # A successfully accepted key frame establishes decoder continuity even
        # when FFmpeg buffers it and returns no image until a later packet.
        self._waiting_for_keyframe = False
        emissions: list[Emission] = []
        for frame in frames:
            decoded = ImageFrame(
                data=frame.data,
                width=frame.width,
                height=frame.height,
                pixel_format=frame.pixel_format,
                codec="none",
                is_key_frame=False,
                intrinsics=payload.intrinsics,
            )
            emissions.append(
                Emission(
                    "frames",
                    envelope.with_payload(
                        decoded,
                        schema_name="ImageFrame",
                        schema_version=1,
                    ),
                )
            )
        self._decoded_frames += len(emissions)
        return tuple(emissions) if emissions else None

    async def stop(self) -> None:
        await self._reset_stream(count_reset=False)
        self._executor = None

    async def health(self) -> HealthReport:
        return HealthReport(
            True,
            "decoded_frames={} dropped_until_keyframe={} resets={} "
            "decode_errors={} waiting_for_keyframe={}".format(
                self._decoded_frames,
                self._dropped_until_keyframe,
                self._resets,
                self._decode_errors,
                self._waiting_for_keyframe,
            ),
        )

    async def _reset_stream(
        self,
        *,
        identity: tuple[str, str | None, str, int, int] | None = None,
        count_reset: bool = True,
    ) -> None:
        had_state = self._stream_identity is not None or self._backend is not None
        if had_state:
            # Always synchronize with the backend lock once a stream exists.
            # If an awaited executor call was cancelled while native decoding
            # continued, checking only ``self._backend`` here could race before
            # the worker publishes the newly-created backend and leak it.
            await run_blocking(self._close_backend_sync, executor=self._executor)
        self._stream_identity = identity
        self._last_sequence = None
        self._waiting_for_keyframe = True
        if count_reset and had_state:
            self._resets += 1

    def _decode_sync(
        self,
        codec: str,
        data: bytes,
    ) -> tuple[DecodedVideoFrame, ...]:
        with self._backend_lock:
            if self._backend is None:
                self._backend = self._backend_factory(codec)
                self._backend_codec = codec
            elif self._backend_codec != codec:
                # The async state machine normally closes it first.  Retain a
                # defensive guard in case this helper is reused independently.
                self._backend.close()
                self._backend = self._backend_factory(codec)
                self._backend_codec = codec
            return tuple(self._backend.decode(data))

    def _close_backend_sync(self) -> None:
        with self._backend_lock:
            backend = self._backend
            self._backend = None
            self._backend_codec = None
            if backend is not None:
                backend.close()


def _normalize_codec(value: str) -> str | None:
    if not isinstance(value, str):
        raise TypeError("image codec must be a string")
    codec = value.strip().lower().lstrip(".")
    if codec in {"", "none", "raw"}:
        return None
    if codec in {"h264", "avc", "avc1"}:
        return "h264"
    if codec in {"h265", "hevc", "hev1", "hvc1"}:
        return "hevc"
    raise ValueError(f"unsupported encoded video codec: {value!r}")


__all__ = [
    "DecodedVideoFrame",
    "VideoDecodeError",
    "VideoDecoderBackend",
    "VideoDecoderBackendFactory",
    "VideoDecoderOperator",
]
