arcly-stream 0.1.4

//! The media frame model — the unit of data that flows through the engine.

use bytes::Bytes;
use serde::{Deserialize, Serialize};
use std::time::Duration;

/// Identifies the codec used to encode a media sample.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[repr(u8)]
#[non_exhaustive]
pub enum CodecId {
    /// H.264 / AVC video.
    H264 = 0,
    /// H.265 / HEVC video.
    H265 = 1,
    /// VP8 video.
    VP8 = 2,
    /// VP9 video.
    VP9 = 3,
    /// AV1 video.
    AV1 = 4,
    /// VVC / H.266 video.
    VVC = 12,
    /// AAC audio.
    AAC = 5,
    /// Opus audio.
    Opus = 6,
    /// G.711 A-law audio.
    G711Alaw = 7,
    /// G.711 µ-law audio.
    G711Ulaw = 8,
    /// G.722 audio.
    G722 = 9,
    /// MP3 audio.
    MP3 = 10,
    /// Raw / passthrough payload (no codec interpretation).
    Raw = 11,
    /// Unknown or unspecified codec.
    Unknown = 255,
}

impl TryFrom<u8> for CodecId {
    type Error = u8;

    fn try_from(v: u8) -> std::result::Result<Self, u8> {
        match v {
            0 => Ok(CodecId::H264),
            1 => Ok(CodecId::H265),
            2 => Ok(CodecId::VP8),
            3 => Ok(CodecId::VP9),
            4 => Ok(CodecId::AV1),
            5 => Ok(CodecId::AAC),
            6 => Ok(CodecId::Opus),
            7 => Ok(CodecId::G711Alaw),
            8 => Ok(CodecId::G711Ulaw),
            9 => Ok(CodecId::G722),
            10 => Ok(CodecId::MP3),
            11 => Ok(CodecId::Raw),
            12 => Ok(CodecId::VVC),
            255 => Ok(CodecId::Unknown),
            n => Err(n),
        }
    }
}

/// Frame classification for video.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum FrameType {
    /// Intra / keyframe (IDR for H.264)
    Key,
    /// Inter / delta frame
    Delta,
    /// Audio frame
    Audio,
}

bitflags::bitflags! {
    /// Bit flags attached to every MediaFrame.
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
    pub struct FrameFlags: u16 {
        /// Frame is the first in the stream.
        const FIRST        = 0b0000_0001;
        /// Frame completes a GOP (Group of Pictures).
        const GOP_END      = 0b0000_0010;
        /// Frame carries sequence/decoder config (SPS/PPS for H.264).
        const CONFIG       = 0b0000_0100;
        /// Frame is a discontinuity marker.
        const DISCONTINUITY = 0b0000_1000;
    }
}

/// A single decoded or encoded media sample.
///
/// `data` is always a [`bytes::Bytes`] slice — zero-copy cloning is safe and
/// cheap.  The same frame can be fanned-out to many subscribers without
/// copying the underlying buffer.
#[derive(Debug, Clone)]
pub struct MediaFrame {
    /// Presentation timestamp in the stream's time base (milliseconds).
    pub pts: i64,
    /// Decode timestamp in the stream's time base (milliseconds).
    pub dts: i64,
    /// Duration of this frame in milliseconds.  `None` = unknown.
    pub duration: Option<u64>,
    /// Encoded / raw payload.
    pub data: Bytes,
    /// Codec that produced this frame.
    pub codec: CodecId,
    /// Frame type classification.
    pub frame_type: FrameType,
    /// Bit flags.
    pub flags: FrameFlags,
    /// Track index (0 = first video, 1 = first audio, etc.).
    pub track_id: u32,
}

impl MediaFrame {
    /// Build a video frame; `is_key` selects [`FrameType::Key`] vs
    /// [`FrameType::Delta`]. Track id defaults to 0.
    pub fn new_video(pts: i64, dts: i64, data: Bytes, codec: CodecId, is_key: bool) -> Self {
        Self {
            pts,
            dts,
            duration: None,
            data,
            codec,
            frame_type: if is_key {
                FrameType::Key
            } else {
                FrameType::Delta
            },
            flags: FrameFlags::empty(),
            track_id: 0,
        }
    }

    /// Build an audio frame (`dts == pts`, [`FrameType::Audio`], track id 1).
    pub fn new_audio(pts: i64, data: Bytes, codec: CodecId) -> Self {
        Self {
            pts,
            dts: pts,
            duration: None,
            data,
            codec,
            frame_type: FrameType::Audio,
            flags: FrameFlags::empty(),
            track_id: 1,
        }
    }

    /// Whether this frame is a video keyframe ([`FrameType::Key`]).
    pub fn is_keyframe(&self) -> bool {
        self.frame_type == FrameType::Key
    }

    /// Whether this frame is audio.
    pub fn is_audio(&self) -> bool {
        self.frame_type == FrameType::Audio
    }

    /// Whether this frame is video (key or delta).
    pub fn is_video(&self) -> bool {
        matches!(self.frame_type, FrameType::Key | FrameType::Delta)
    }

    /// The presentation timestamp as a [`Duration`] (from the absolute `pts`).
    pub fn pts_duration(&self) -> Duration {
        Duration::from_millis(self.pts.unsigned_abs())
    }
}

/// Video-specific frame carrying extra metadata.
#[derive(Debug, Clone)]
pub struct VideoFrame {
    /// The underlying media frame.
    pub inner: MediaFrame,
    /// Coded width in pixels.
    pub width: u32,
    /// Coded height in pixels.
    pub height: u32,
    /// Frame-rate numerator.
    pub fps_num: u32,
    /// Frame-rate denominator.
    pub fps_den: u32,
}

/// Audio-specific frame carrying extra metadata.
#[derive(Debug, Clone)]
pub struct AudioFrame {
    /// The underlying media frame.
    pub inner: MediaFrame,
    /// Sample rate in Hz.
    pub sample_rate: u32,
    /// Channel count.
    pub channels: u8,
    /// Bits per sample.
    pub bits_per_sample: u8,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn video_keyframe_classification() {
        let key = MediaFrame::new_video(0, 0, Bytes::from_static(b"x"), CodecId::H264, true);
        assert!(key.is_keyframe());
        assert!(key.is_video());
        assert!(!key.is_audio());
        assert_eq!(key.frame_type, FrameType::Key);

        let delta = MediaFrame::new_video(1, 1, Bytes::new(), CodecId::H264, false);
        assert!(!delta.is_keyframe());
        assert!(delta.is_video());
    }

    #[test]
    fn audio_frame_defaults_track_and_type() {
        let a = MediaFrame::new_audio(10, Bytes::new(), CodecId::AAC);
        assert!(a.is_audio());
        assert!(!a.is_video());
        assert_eq!(a.track_id, 1);
        assert_eq!(a.dts, a.pts); // audio has no separate decode order
    }

    #[test]
    fn codec_id_roundtrips_through_u8() {
        for id in [CodecId::H264, CodecId::Opus, CodecId::Raw, CodecId::Unknown] {
            assert_eq!(CodecId::try_from(id as u8), Ok(id));
        }
        assert_eq!(CodecId::try_from(200u8), Err(200));
    }

    #[test]
    fn frame_flags_compose() {
        let flags = FrameFlags::CONFIG | FrameFlags::FIRST;
        assert!(flags.contains(FrameFlags::CONFIG));
        assert!(!flags.contains(FrameFlags::GOP_END));
    }
}