yovision/Brain/src/yovision_brain/decode/y4m.py

"""Minimal streaming YUV4MPEG2 decoder for anonymous local fixtures.

The backend intentionally supports only uncompressed C444 streams. Production
codecs and RTSP belong behind the same decoder port in later tasks.
"""

from __future__ import annotations

from collections.abc import Iterable, Iterator
from dataclasses import dataclass

from yovision_brain.input import CancellationToken, InputPacket

from .models import DecodedFrame, DecoderError

_MAX_HEADER_BYTES = 4096
_MAX_FRAME_BYTES = 256 * 1024 * 1024


class _Cancelled(Exception):
    pass


class _PacketReader:
    def __init__(
        self,
        packets: Iterable[InputPacket],
        cancellation: CancellationToken | None,
    ) -> None:
        self._packets = iter(packets)
        self._cancellation = cancellation
        self._buffer = bytearray()
        self._ended = False
        self.first_packet: InputPacket | None = None

    def _fill(self) -> bool:
        if self._cancellation is not None and self._cancellation.cancelled:
            raise _Cancelled
        if self._ended:
            return False
        try:
            packet = next(self._packets)
        except StopIteration:
            self._ended = True
            return False
        if packet.media_format != "container-bytes":
            raise DecoderError(f"Y4M decoder does not support {packet.media_format!r}")
        if self.first_packet is None:
            self.first_packet = packet
        else:
            first = self.first_packet
            if (packet.logical_device_id, packet.profile_id) != (
                first.logical_device_id,
                first.profile_id,
            ):
                raise DecoderError("input identity changed inside one local video stream")
        self._buffer.extend(packet.payload)
        return True

    def line(self, *, allow_clean_eof: bool = False) -> bytes | None:
        while True:
            newline = self._buffer.find(b"\n")
            if newline >= 0:
                result = bytes(self._buffer[:newline])
                del self._buffer[: newline + 1]
                return result
            if len(self._buffer) > _MAX_HEADER_BYTES:
                raise DecoderError("Y4M header exceeds the safe size limit")
            if not self._fill():
                if not self._buffer and allow_clean_eof:
                    return None
                raise DecoderError("truncated Y4M header")

    def exact(self, size: int) -> bytes:
        while len(self._buffer) < size:
            if not self._fill():
                raise DecoderError("truncated Y4M frame payload")
        result = bytes(self._buffer[:size])
        del self._buffer[:size]
        return result


@dataclass(frozen=True, slots=True)
class _Header:
    width: int
    height: int
    fps_numerator: int
    fps_denominator: int


def _positive_int(value: bytes, field: str) -> int:
    try:
        result = int(value)
    except ValueError as exc:
        raise DecoderError(f"invalid Y4M {field}") from exc
    if result <= 0:
        raise DecoderError(f"invalid Y4M {field}")
    return result


def _parse_header(line: bytes) -> _Header:
    parts = line.split()
    if not parts or parts[0] != b"YUV4MPEG2":
        raise DecoderError("unsupported local video format; expected YUV4MPEG2")
    fields = {part[:1]: part[1:] for part in parts[1:] if len(part) > 1}
    if fields.get(b"C", b"444") not in {b"444", b"444jpeg"}:
        raise DecoderError("unsupported Y4M chroma; only C444 is supported")
    width = _positive_int(fields.get(b"W", b""), "width")
    height = _positive_int(fields.get(b"H", b""), "height")
    fps_parts = fields.get(b"F", b"").split(b":", 1)
    if len(fps_parts) != 2:
        raise DecoderError("invalid Y4M frame rate")
    header = _Header(
        width=width,
        height=height,
        fps_numerator=_positive_int(fps_parts[0], "frame rate numerator"),
        fps_denominator=_positive_int(fps_parts[1], "frame rate denominator"),
    )
    if header.width * header.height * 3 > _MAX_FRAME_BYTES:
        raise DecoderError("Y4M frame exceeds the safe size limit")
    return header


class Y4MDecoder:
    media_formats = frozenset({"container-bytes"})

    def decode(
        self,
        packets: Iterable[InputPacket],
        cancellation: CancellationToken | None = None,
    ) -> Iterator[DecodedFrame]:
        reader = _PacketReader(packets, cancellation)
        try:
            header_line = reader.line()
            assert header_line is not None
            header = _parse_header(header_line)
            first = reader.first_packet
            if first is None:
                raise DecoderError("local video input is empty")
            if (first.width, first.height) != (header.width, header.height):
                raise DecoderError(
                    "Y4M dimensions do not match the configured input profile "
                    f"({header.width}x{header.height} != {first.width}x{first.height})"
                )
            interval_ns = round(1_000_000_000 * header.fps_denominator / header.fps_numerator)
            frame_size = header.width * header.height * 3
            sequence = 0
            while True:
                frame_header = reader.line(allow_clean_eof=True)
                if frame_header is None:
                    return
                if frame_header != b"FRAME":
                    raise DecoderError(f"invalid Y4M frame header at frame {sequence}")
                payload = reader.exact(frame_size)
                yield DecodedFrame(
                    sequence=sequence,
                    timestamp_ns=sequence * interval_ns,
                    logical_device_id=first.logical_device_id,
                    profile_id=first.profile_id,
                    width=header.width,
                    height=header.height,
                    pixel_format="yuv444p",
                    payload=payload,
                )
                sequence += 1
        except _Cancelled:
            return