"""Bounded media/runtime primitives for the LAB E15 shadow inference worker."""

from __future__ import annotations

import threading
import time
from collections import deque
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any

import numpy as np


class ShadowRuntimeError(RuntimeError):
    pass


@dataclass(frozen=True, slots=True)
class CameraFragmentMetadata:
    ingress_sequence: int
    source_sequence: int
    captured_at_epoch_ns: int
    worker_received_monotonic: float


@dataclass(frozen=True, slots=True)
class DecodedCameraFrame:
    frame_index: int
    metadata: CameraFragmentMetadata
    image: np.ndarray
    decoded_monotonic: float
    decode_age_ms: float


class IncrementalMediaBuffer:
    """A bounded non-seekable blocking byte source accepted by PyAV."""

    def __init__(self, maximum_buffer_bytes: int) -> None:
        if maximum_buffer_bytes < 1024:
            raise ValueError("media buffer bound is too small")
        self._maximum_buffer_bytes = maximum_buffer_bytes
        self._buffer = bytearray()
        self._condition = threading.Condition()
        self._finished = False
        self._failure: BaseException | None = None
        self._bytes_published = 0
        self._bytes_read = 0
        self._maximum_depth = 0

    def readable(self) -> bool:
        return True

    def read(self, size: int = -1) -> bytes:
        with self._condition:
            self._condition.wait_for(
                lambda: bool(self._buffer) or self._finished or self._failure is not None
            )
            if self._failure is not None:
                raise ShadowRuntimeError("incremental media source failed") from self._failure
            if not self._buffer:
                return b""
            count = len(self._buffer) if size < 0 else min(size, len(self._buffer))
            result = bytes(self._buffer[:count])
            del self._buffer[:count]
            self._bytes_read += count
            self._condition.notify_all()
            return result

    def append(self, payload: bytes) -> None:
        if not payload:
            raise ShadowRuntimeError("empty fMP4 payload is not admitted")
        with self._condition:
            if self._finished or self._failure is not None:
                raise ShadowRuntimeError("incremental media source is closed")
            if len(self._buffer) + len(payload) > self._maximum_buffer_bytes:
                failure = ShadowRuntimeError("incremental media buffer exceeded its hard bound")
                self._failure = failure
                self._condition.notify_all()
                raise failure
            self._buffer.extend(payload)
            self._bytes_published += len(payload)
            self._maximum_depth = max(self._maximum_depth, len(self._buffer))
            self._condition.notify_all()

    def finish(self) -> None:
        with self._condition:
            self._finished = True
            self._condition.notify_all()

    def fail(self, failure: BaseException) -> None:
        with self._condition:
            if self._failure is None:
                self._failure = failure
            self._condition.notify_all()

    def snapshot(self) -> dict[str, int | bool]:
        with self._condition:
            return {
                "maximum_buffer_bytes": self._maximum_buffer_bytes,
                "depth_bytes": len(self._buffer),
                "maximum_depth_bytes": self._maximum_depth,
                "bytes_published": self._bytes_published,
                "bytes_read": self._bytes_read,
                "finished": self._finished,
                "failed": self._failure is not None,
            }


class CameraMetadataQueue:
    """Fail-closed segment metadata queue; dropping a fragment would corrupt H.264."""

    def __init__(self, capacity: int) -> None:
        if capacity < 2:
            raise ValueError("camera metadata capacity is too small")
        self._capacity = capacity
        self._items: deque[CameraFragmentMetadata] = deque()
        self._condition = threading.Condition()
        self._finished = False
        self._published = 0
        self._consumed = 0
        self._maximum_depth = 0

    def publish(self, value: CameraFragmentMetadata) -> None:
        with self._condition:
            if self._finished:
                raise ShadowRuntimeError("camera metadata queue is closed")
            if len(self._items) >= self._capacity:
                raise ShadowRuntimeError("camera metadata queue exceeded its hard bound")
            self._items.append(value)
            self._published += 1
            self._maximum_depth = max(self._maximum_depth, len(self._items))
            self._condition.notify_all()

    def take(self) -> CameraFragmentMetadata:
        with self._condition:
            self._condition.wait_for(lambda: bool(self._items) or self._finished)
            if not self._items:
                raise ShadowRuntimeError("decoder emitted a frame without segment metadata")
            self._consumed += 1
            return self._items.popleft()

    def finish(self) -> None:
        with self._condition:
            self._finished = True
            self._condition.notify_all()

    def snapshot(self) -> dict[str, int | bool]:
        with self._condition:
            return {
                "capacity": self._capacity,
                "depth": len(self._items),
                "maximum_depth": self._maximum_depth,
                "published": self._published,
                "consumed": self._consumed,
                "finished": self._finished,
            }


class PersistentFmp4Decoder:
    """Decode one committed fMP4 epoch incrementally without writing frames."""

    def __init__(
        self,
        *,
        on_frame: Callable[[DecodedCameraFrame], None],
        width: int = 800,
        height: int = 600,
        maximum_buffer_bytes: int = 8 * 1024 * 1024,
        metadata_capacity: int = 16,
    ) -> None:
        if width < 1 or height < 1:
            raise ValueError("decoder resolution is invalid")
        self._on_frame = on_frame
        self._width = width
        self._height = height
        self._media = IncrementalMediaBuffer(maximum_buffer_bytes)
        self._metadata = CameraMetadataQueue(metadata_capacity)
        self._thread = threading.Thread(
            target=self._decode,
            name="lab-e15-fmp4-decoder",
            daemon=True,
        )
        self._started = False
        self._init_seen = False
        self._last_source_sequence: int | None = None
        self._source_sequence_gaps = 0
        self._decoded_frames = 0
        self._failure: BaseException | None = None

    def start(self) -> None:
        if self._started:
            raise ShadowRuntimeError("fMP4 decoder was already started")
        self._started = True
        self._thread.start()

    def feed_init(self, payload: bytes) -> None:
        if not self._started or self._init_seen or self._last_source_sequence is not None:
            raise ShadowRuntimeError("camera init ordering is invalid")
        self._init_seen = True
        self._media.append(payload)

    def feed_segment(self, metadata: CameraFragmentMetadata, payload: bytes) -> None:
        if not self._init_seen:
            raise ShadowRuntimeError("camera media arrived before init")
        if metadata.source_sequence < 1:
            raise ShadowRuntimeError("camera source sequence must be positive")
        if (
            self._last_source_sequence is not None
            and metadata.source_sequence <= self._last_source_sequence
        ):
            raise ShadowRuntimeError(
                "camera source sequence is not increasing: "
                f"previous {self._last_source_sequence}, got {metadata.source_sequence}"
            )
        expected = 1 if self._last_source_sequence is None else self._last_source_sequence + 1
        self._source_sequence_gaps += metadata.source_sequence - expected
        self._metadata.publish(metadata)
        try:
            self._media.append(payload)
        except BaseException as exc:
            self._media.fail(exc)
            self._metadata.finish()
            raise
        self._last_source_sequence = metadata.source_sequence

    def finish_input(self) -> None:
        self._media.finish()
        self._metadata.finish()

    def join(self, timeout_seconds: float = 30.0) -> None:
        self._thread.join(timeout=timeout_seconds)
        if self._thread.is_alive():
            raise ShadowRuntimeError("persistent fMP4 decoder did not stop")
        if self._failure is not None:
            raise ShadowRuntimeError("persistent fMP4 decoder failed") from self._failure
        metadata = self._metadata.snapshot()
        if metadata["published"] != metadata["consumed"] or metadata["depth"] != 0:
            raise ShadowRuntimeError("camera segment/frame accounting differs")

    def snapshot(self) -> dict[str, Any]:
        return {
            "init_seen": self._init_seen,
            "last_source_sequence": self._last_source_sequence,
            "source_sequence_gaps": self._source_sequence_gaps,
            "decoded_frames": self._decoded_frames,
            "failed": self._failure is not None,
            "media": self._media.snapshot(),
            "metadata": self._metadata.snapshot(),
        }

    def _decode(self) -> None:
        try:
            import av

            container = av.open(
                self._media,
                mode="r",
                format="mp4",
                options={"probesize": "32", "analyzeduration": "0"},
            )
            try:
                for frame in container.decode(video=0):
                    metadata = self._metadata.take()
                    image = frame.to_ndarray(format="rgb24")
                    if image.shape != (self._height, self._width, 3):
                        raise ShadowRuntimeError("decoded camera resolution changed")
                    image.setflags(write=False)
                    decoded_monotonic = time.perf_counter()
                    self._on_frame(
                        DecodedCameraFrame(
                            frame_index=self._decoded_frames,
                            metadata=metadata,
                            image=image,
                            decoded_monotonic=decoded_monotonic,
                            decode_age_ms=max(
                                0.0,
                                (
                                    decoded_monotonic
                                    - metadata.worker_received_monotonic
                                )
                                * 1000,
                            ),
                        )
                    )
                    self._decoded_frames += 1
            finally:
                container.close()
        except BaseException as exc:
            self._failure = exc
            self._media.fail(exc)
            self._metadata.finish()
