"""One-sample fMP4 → persistent H.264 decoder, no recording/file dependency.

Reuses the archive's bounded ISO-BMFF timing validator, not its materializer.
Only the supplied init and current complete fragment are inspected. Codec state
persists across fragments; no later fragment or source EOF is supplied to make
the current frame appear. Other codecs/reordered/multi-sample streams fail
closed and require a different explicit profile. Run inside a supervised,
memory-limited CPU child; native decoder allocations are not Python byte counts.
"""

from __future__ import annotations

import importlib
import io
from fractions import Fraction
from typing import Any

from k1link.media_fragments import (
    ParseBudget,
    single_sample_payload,
    video_fragment_timing,
    video_timing,
)


class StreamingDecodeError(ValueError):
    pass


class FragmentDecoder:
    def __init__(self) -> None:
        self._av = importlib.import_module("av")
        if self._av.__version__ != "18.0.0":
            raise StreamingDecodeError("unqualified PyAV version")
        # Import before readiness, not on the first source frame's deadline.
        importlib.import_module("numpy")
        self._codec: Any = None
        self._init: bytes | None = None
        self._timing: Any = None
        self._next_dts: int | None = None
        self.frames = 0
        self.failed = False

    def configure(self, payload: bytes) -> None:
        if self._init is not None or self.failed or not 0 < len(payload) <= 65536:
            self.failed = True
            raise StreamingDecodeError("invalid or repeated codec initialization")
        try:
            timing = video_timing(payload, ParseBudget(128, 1))
            with self._av.open(io.BytesIO(payload), format="mp4", mode="r") as container:
                if len(container.streams) != 1 or len(container.streams.video) != 1:
                    raise StreamingDecodeError("one video track is required")
                description = container.streams.video[0].codec_context
                if description.name != "h264" or (description.width, description.height) != (
                    800,
                    600,
                ):
                    raise StreamingDecodeError("codec or calibrated native raster changed")
                extra = description.extradata
                if not extra or len(extra) > 65536:
                    raise StreamingDecodeError("bounded H264 codec configuration required")
                codec = self._av.CodecContext.create("h264", "r")
                codec.extradata = extra
                codec.thread_count = 1
                codec.thread_type = "SLICE"
                codec.open()
            self._codec, self._init, self._timing = codec, bytes(payload), timing
        except Exception:
            self.failed = True
            raise

    def decode(self, payload: bytes) -> Any:
        if self.failed or self._init is None or not 0 < len(payload) <= 1024 * 1024:
            self.failed = True
            raise StreamingDecodeError("uninitialized, failed or oversized fragment")
        try:
            timing = video_fragment_timing(payload, self._timing, ParseBudget(128, 1))
            if self.frames == 0 and not timing.random_access:
                raise StreamingDecodeError("epoch must start at a random-access fragment")
            if self._next_dts is not None and timing.base_decode_time != self._next_dts:
                raise StreamingDecodeError("fragment codec timeline is discontinuous")
            # Opening a new demuxer per fragment repeats FFmpeg stream probing.
            # Validate the narrow explicit layout, then feed the one AVCC sample
            # to the persistent codec. No future packet or flush is supplied.
            packet = self._av.Packet(single_sample_payload(payload, self._timing))
            packet.pts = packet.dts = timing.base_decode_time
            packet.duration = timing.duration_units
            packet.time_base = Fraction(1, self._timing.timescale)
            frames = self._codec.decode(packet)
            if len(frames) != 1 or self._codec.has_b_frames:
                raise StreamingDecodeError("current frame requires future input or reordering")
            frame = frames[0]
            if (frame.width, frame.height) != (800, 600) or frame.pts != packet.pts:
                raise StreamingDecodeError("decoded raster or frame identity changed")
            if self.frames == 0 and not frame.key_frame:
                raise StreamingDecodeError("decoder did not confirm the initial keyframe")
            image = frame.to_ndarray(format="bgr24")
            if image.shape != (600, 800, 3) or image.nbytes != 1_440_000:
                raise StreamingDecodeError("invalid native decoder output")
            image.setflags(write=False)
            self._next_dts = timing.base_decode_time + timing.duration_units
            self.frames += 1
            return image
        except Exception:
            self.failed = True
            raise

    def close(self) -> None:
        self.failed = True
        self._codec = self._init = self._timing = None