arcly-stream 0.1.0

//! 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 {
    // Video
    H264 = 0,
    H265 = 1,
    VP8 = 2,
    VP9 = 3,
    AV1 = 4,
    /// VVC / H.266.
    VVC = 12,
    // Audio
    AAC = 5,
    Opus = 6,
    G711Alaw = 7,
    G711Ulaw = 8,
    G722 = 9,
    MP3 = 10,
    // Raw / passthrough
    Raw = 11,
    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 {
    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,
        }
    }

    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,
        }
    }

    pub fn is_keyframe(&self) -> bool {
        self.frame_type == FrameType::Key
    }

    pub fn is_audio(&self) -> bool {
        self.frame_type == FrameType::Audio
    }

    pub fn is_video(&self) -> bool {
        matches!(self.frame_type, FrameType::Key | FrameType::Delta)
    }

    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 {
    pub inner: MediaFrame,
    pub width: u32,
    pub height: u32,
    pub fps_num: u32,
    pub fps_den: u32,
}

/// Audio-specific frame carrying extra metadata.
#[derive(Debug, Clone)]
pub struct AudioFrame {
    pub inner: MediaFrame,
    pub sample_rate: u32,
    pub channels: u8,
    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));
    }
}