arcly-stream 0.8.1

//! FLV audio/video tag bodies as carried inside RTMP message payloads.
//!
//! Handles the two codecs every RTMP deployment uses — **H.264/AVC** video and
//! **AAC** audio — in both directions:
//!
//! * **Ingest:** decode the AVC decoder-configuration record and AudioSpecificConfig,
//!   turn AVCC NAL units into Annex-B and raw AAC into ADTS, so the engine stores
//!   frames in a player-ready elementary form (what [`MpegTsMuxer`] consumes).
//! * **Egress (play):** rebuild the configuration records and FLV tag bodies from
//!   the engine's Annex-B / ADTS frames.
//!
//! [`MpegTsMuxer`]: crate::packager::MpegTsMuxer

use bytes::{Buf, BufMut, Bytes, BytesMut};

/// FLV codec id for AVC (H.264) in a video tag.
pub const VIDEO_CODEC_AVC: u8 = 7;
/// FLV sound format for AAC in an audio tag.
pub const AUDIO_FORMAT_AAC: u8 = 10;

/// FLV `FrameType` nibble: a keyframe (IDR).
const FRAME_TYPE_KEY: u8 = 1;
/// FLV `FrameType` nibble: an inter (delta) frame.
const FRAME_TYPE_INTER: u8 = 2;

/// `AVCPacketType` / `AACPacketType`: sequence (decoder-config) header.
pub const PKT_SEQUENCE_HEADER: u8 = 0;
/// `AVCPacketType` / `AACPacketType`: raw NALU / raw AAC.
pub const PKT_RAW: u8 = 1;

// ── Video: AVCDecoderConfigurationRecord ────────────────────────────────────

/// The `avcC` parser now lives in the codec layer so non-RTMP consumers (the
/// transcode pipeline) can reuse it; re-exported here for the FLV path's callers.
pub use crate::codec::h264::AvcConfig;

/// A parsed video tag header: codec, frame type, AVC packet type, and the
/// composition-time offset (PTS − DTS, in milliseconds).
pub struct VideoTagHeader {
    pub is_keyframe: bool,
    pub avc_packet_type: u8,
    pub composition_time: i32,
    /// Offset of the body following the 5-byte AVC header.
    pub body_offset: usize,
}

/// Parse the fixed video-tag header. Returns `None` for non-AVC tags.
pub fn parse_video_header(payload: &[u8]) -> Option<VideoTagHeader> {
    let b0 = *payload.first()?;
    let frame_type = b0 >> 4;
    let codec_id = b0 & 0x0F;
    if codec_id != VIDEO_CODEC_AVC {
        return None;
    }
    let avc_packet_type = *payload.get(1)?;
    // 24-bit signed composition time offset (big-endian), then sign-extended.
    let raw = payload.get(2..5)?.get_uint(3) as i32;
    let composition_time = (raw << 8) >> 8; // sign-extend from 24 bits
    Some(VideoTagHeader {
        is_keyframe: frame_type == FRAME_TYPE_KEY,
        avc_packet_type,
        composition_time,
        body_offset: 5,
    })
}

/// Build a video tag body for egress: `[frametype|codec][avc_pkt_type][cts:24][body]`.
pub fn build_video_tag(is_keyframe: bool, avc_packet_type: u8, cts: i32, body: &[u8]) -> Bytes {
    let mut out = BytesMut::with_capacity(body.len() + 5);
    let frame_type = if is_keyframe {
        FRAME_TYPE_KEY
    } else {
        FRAME_TYPE_INTER
    };
    out.put_u8((frame_type << 4) | VIDEO_CODEC_AVC);
    out.put_u8(avc_packet_type);
    out.put_uint((cts & 0x00FF_FFFF) as u64, 3); // 24-bit big-endian CTS
    out.put_slice(body);
    out.freeze()
}

// ── Audio: AudioSpecificConfig ↔ ADTS ───────────────────────────────────────

/// Parsed AAC `AudioSpecificConfig` (the audio sequence header).
#[derive(Debug, Clone, Copy, Default)]
pub struct AudioConfig {
    /// MPEG-4 audio object type (2 = AAC-LC).
    pub object_type: u8,
    /// MPEG-4 sampling-frequency index (4 = 44.1 kHz, 3 = 48 kHz, …).
    pub freq_index: u8,
    /// Channel configuration (1 = mono, 2 = stereo, …).
    pub channels: u8,
}

impl AudioConfig {
    /// Parse a 2-byte (or longer) AudioSpecificConfig.
    pub fn parse(asc: &[u8]) -> Option<Self> {
        if asc.len() < 2 {
            return None;
        }
        let object_type = asc[0] >> 3;
        let freq_index = ((asc[0] & 0x07) << 1) | (asc[1] >> 7);
        let channels = (asc[1] >> 3) & 0x0F;
        Some(Self {
            object_type,
            freq_index,
            channels,
        })
    }

    /// Reconstruct the 2-byte AudioSpecificConfig.
    pub fn to_asc(self) -> Bytes {
        let b0 = (self.object_type << 3) | (self.freq_index >> 1);
        let b1 = ((self.freq_index & 0x01) << 7) | ((self.channels & 0x0F) << 3);
        Bytes::copy_from_slice(&[b0, b1])
    }

    /// Prefix raw AAC with a 7-byte ADTS header (what MPEG-TS carries).
    pub fn to_adts(self, raw: &[u8]) -> Bytes {
        let frame_len = raw.len() + 7;
        let profile = self.object_type.saturating_sub(1); // ADTS profile = AOT − 1
        let hdr = [
            0xFF,
            0xF1, // MPEG-4, layer 0, no CRC
            (profile << 6) | (self.freq_index << 2) | ((self.channels >> 2) & 0x01),
            ((self.channels & 0x03) << 6) | ((frame_len >> 11) & 0x03) as u8,
            ((frame_len >> 3) & 0xFF) as u8,
            (((frame_len & 0x07) << 5) as u8) | 0x1F,
            0xFC,
        ];
        let mut out = BytesMut::with_capacity(frame_len);
        out.put_slice(&hdr);
        out.put_slice(raw);
        out.freeze()
    }

    /// Derive the config from a 7-byte ADTS header (used on the egress path,
    /// where the engine stored ADTS frames).
    pub fn from_adts(hdr: &[u8]) -> Option<Self> {
        if hdr.len() < 7 || hdr[0] != 0xFF || (hdr[1] & 0xF0) != 0xF0 {
            return None;
        }
        let profile = (hdr[2] >> 6) & 0x03;
        let freq_index = (hdr[2] >> 2) & 0x0F;
        let channels = ((hdr[2] & 0x01) << 2) | (hdr[3] >> 6);
        Some(Self {
            object_type: profile + 1,
            freq_index,
            channels,
        })
    }
}

/// Build an audio tag body for egress: `[soundfmt|rate|size|type][aac_pkt][body]`.
/// The sound-rate/size/type nibble is fixed to AAC's conventional `0xF` (44 kHz,
/// 16-bit, stereo) — decoders read the real values from the AudioSpecificConfig.
pub fn build_audio_tag(aac_packet_type: u8, body: &[u8]) -> Bytes {
    let mut out = BytesMut::with_capacity(body.len() + 2);
    out.put_u8((AUDIO_FORMAT_AAC << 4) | 0x0F);
    out.put_u8(aac_packet_type);
    out.put_slice(body);
    out.freeze()
}

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

    #[test]
    fn avc_config_roundtrips_through_annexb_and_record() {
        let sps = vec![0x67, 0x42, 0x00, 0x1F, 0xAA];
        let pps = vec![0x68, 0xCE, 0x3C, 0x80];
        let rec = {
            let mut r = vec![1, sps[1], sps[2], sps[3], 0xFF, 0xE1];
            r.extend_from_slice(&(sps.len() as u16).to_be_bytes());
            r.extend_from_slice(&sps);
            r.push(1);
            r.extend_from_slice(&(pps.len() as u16).to_be_bytes());
            r.extend_from_slice(&pps);
            r
        };
        let cfg = AvcConfig::parse(&rec).expect("parse");
        assert_eq!(cfg.nal_length_size, 4);
        assert_eq!(cfg.sps, vec![sps.clone()]);
        assert_eq!(cfg.pps, vec![pps.clone()]);

        // Annex-B form starts with a 4-byte start code then the SPS.
        let annexb = cfg.to_annexb();
        assert_eq!(&annexb[..5], &[0, 0, 0, 1, 0x67]);

        // Rebuilt record parses back to the same parameter sets.
        let rebuilt = AvcConfig::parse(&cfg.to_avc_record()).expect("reparse");
        assert_eq!(rebuilt.sps, cfg.sps);
        assert_eq!(rebuilt.pps, cfg.pps);
    }

    #[test]
    fn avcc_nalus_become_annexb_with_optional_params() {
        let cfg = AvcConfig {
            nal_length_size: 4,
            sps: vec![vec![0x67, 0x42]],
            pps: vec![vec![0x68, 0xCE]],
        };
        // One AVCC NAL: length 2 + [0x65, 0xAA].
        let avcc = [0, 0, 0, 2, 0x65, 0xAA];
        let plain = cfg.avcc_to_annexb(&avcc, false).unwrap();
        assert_eq!(&plain[..], &[0, 0, 0, 1, 0x65, 0xAA]);
        let with_params = cfg.avcc_to_annexb(&avcc, true).unwrap();
        // SPS + PPS + slice, each start-code prefixed.
        assert_eq!(&with_params[..5], &[0, 0, 0, 1, 0x67]);
        assert!(with_params.ends_with(&[0x65, 0xAA]));
    }

    #[test]
    fn audio_config_asc_and_adts_roundtrip() {
        // AAC-LC (AOT 2), 44.1 kHz (index 4), stereo.
        let cfg = AudioConfig {
            object_type: 2,
            freq_index: 4,
            channels: 2,
        };
        let asc = cfg.to_asc();
        let reparsed = AudioConfig::parse(&asc).unwrap();
        assert_eq!(reparsed.object_type, 2);
        assert_eq!(reparsed.freq_index, 4);
        assert_eq!(reparsed.channels, 2);

        let adts = cfg.to_adts(&[0x11, 0x22, 0x33]);
        assert_eq!(adts.len(), 10); // 7 header + 3 payload
        assert_eq!(adts[0], 0xFF);
        let from_hdr = AudioConfig::from_adts(&adts).unwrap();
        assert_eq!(from_hdr.freq_index, 4);
        assert_eq!(from_hdr.channels, 2);
        assert_eq!(from_hdr.object_type, 2);
    }

    #[test]
    fn video_header_parses_keyframe_and_cts() {
        // keyframe (1) + AVC (7), pkt type 1, cts = 0x000040 (64).
        let tag = build_video_tag(true, PKT_RAW, 64, &[0xAB]);
        let h = parse_video_header(&tag).unwrap();
        assert!(h.is_keyframe);
        assert_eq!(h.avc_packet_type, PKT_RAW);
        assert_eq!(h.composition_time, 64);
        assert_eq!(&tag[h.body_offset..], &[0xAB]);
    }

    #[test]
    fn negative_composition_time_sign_extends() {
        let tag = build_video_tag(false, PKT_RAW, -10, &[]);
        let h = parse_video_header(&tag).unwrap();
        assert_eq!(h.composition_time, -10);
    }
}