lvqr-ingest 1.1.0

//! Fragmented MP4 (CMAF) box writer.
//!
//! Generates fMP4 init segments (ftyp + moov) and media segments (moof + mdat)
//! compatible with MSE SourceBuffer and the MoQ ecosystem.
//!
//! Written manually with `bytes::BytesMut` -- no external MP4 crate needed.

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

use super::flv::{AudioConfig, VideoConfig};

// --- Box writing helpers ---

fn write_box(buf: &mut BytesMut, box_type: &[u8; 4], f: impl FnOnce(&mut BytesMut)) {
    let start = buf.len();
    buf.put_u32(0); // placeholder for size
    buf.put_slice(box_type);
    f(buf);
    let size = (buf.len() - start) as u32;
    let start_bytes = size.to_be_bytes();
    buf[start..start + 4].copy_from_slice(&start_bytes);
}

fn write_full_box(buf: &mut BytesMut, box_type: &[u8; 4], version: u8, flags: u32, f: impl FnOnce(&mut BytesMut)) {
    write_box(buf, box_type, |buf| {
        buf.put_u8(version);
        buf.put_u8((flags >> 16) as u8);
        buf.put_u8((flags >> 8) as u8);
        buf.put_u8(flags as u8);
        f(buf);
    });
}

/// Write an MPEG-4 descriptor (`ISO/IEC 14496-1` §8.3.3) with a variable-length
/// size prefix.
///
/// The size is encoded as 1-4 bytes, each holding 7 payload bits with the
/// MSB set on every byte except the last. Fixed 4-byte form is used here
/// rather than minimal encoding so the descriptor header width is the same
/// regardless of payload size, which keeps `patch`-style callers simple.
/// The 4-byte form supports payloads up to 2^28 - 1 bytes, which is far
/// more than any realistic AudioSpecificConfig or DecoderConfigDescriptor.
fn write_mpeg4_descriptor(buf: &mut BytesMut, tag: u8, f: impl FnOnce(&mut BytesMut)) {
    buf.put_u8(tag);
    let size_pos = buf.len();
    buf.put_bytes(0, 4); // 4-byte placeholder for the length field
    let payload_start = buf.len();
    f(buf);
    let size = buf.len() - payload_start;
    assert!(size < (1 << 28), "MPEG-4 descriptor payload exceeds 28-bit size field");
    buf[size_pos] = 0x80 | ((size >> 21) & 0x7F) as u8;
    buf[size_pos + 1] = 0x80 | ((size >> 14) & 0x7F) as u8;
    buf[size_pos + 2] = 0x80 | ((size >> 7) & 0x7F) as u8;
    buf[size_pos + 3] = (size & 0x7F) as u8;
}

// --- Init segment generation ---

/// Generate a video init segment (ftyp + moov) for H.264/AVC.
/// If width/height are 0, MSE may reject the segment -- pass actual resolution when known.
pub fn video_init_segment(config: &VideoConfig) -> Bytes {
    video_init_segment_with_size(config, 0, 0)
}

/// Generate a video init segment with explicit dimensions.
pub fn video_init_segment_with_size(config: &VideoConfig, width: u16, height: u16) -> Bytes {
    let mut buf = BytesMut::with_capacity(512);

    // ftyp
    write_box(&mut buf, b"ftyp", |buf| {
        buf.put_slice(b"isom"); // major_brand
        buf.put_u32(0); // minor_version
        buf.put_slice(b"isom");
        buf.put_slice(b"iso6");
        buf.put_slice(b"msdh");
        buf.put_slice(b"msix");
    });

    // moov
    write_box(&mut buf, b"moov", |buf| {
        // mvhd
        write_full_box(buf, b"mvhd", 0, 0, |buf| {
            buf.put_u32(0); // creation_time
            buf.put_u32(0); // modification_time
            buf.put_u32(90000); // timescale
            buf.put_u32(0); // duration
            buf.put_u32(0x00010000); // rate = 1.0
            buf.put_u16(0x0100); // volume = 1.0
            buf.put_bytes(0, 10); // reserved
            // identity matrix (9 x u32)
            for &v in &[0x00010000u32, 0, 0, 0, 0x00010000, 0, 0, 0, 0x40000000] {
                buf.put_u32(v);
            }
            buf.put_bytes(0, 24); // pre_defined
            buf.put_u32(2); // next_track_ID
        });

        // trak
        write_box(buf, b"trak", |buf| {
            // tkhd
            write_full_box(buf, b"tkhd", 0, 0x03, |buf| {
                buf.put_u32(0); // creation_time
                buf.put_u32(0); // modification_time
                buf.put_u32(1); // track_ID
                buf.put_u32(0); // reserved
                buf.put_u32(0); // duration
                buf.put_bytes(0, 8); // reserved
                buf.put_u16(0); // layer
                buf.put_u16(0); // alternate_group
                buf.put_u16(0); // volume (0 for video)
                buf.put_u16(0); // reserved
                // identity matrix
                for &v in &[0x00010000u32, 0, 0, 0, 0x00010000, 0, 0, 0, 0x40000000] {
                    buf.put_u32(v);
                }
                buf.put_u32((width as u32) << 16); // width (fixed-point 16.16)
                buf.put_u32((height as u32) << 16); // height (fixed-point 16.16)
            });

            // mdia
            write_box(buf, b"mdia", |buf| {
                // mdhd
                write_full_box(buf, b"mdhd", 0, 0, |buf| {
                    buf.put_u32(0); // creation_time
                    buf.put_u32(0); // modification_time
                    buf.put_u32(90000); // timescale
                    buf.put_u32(0); // duration
                    buf.put_u32(0x55C40000); // language = "und"
                });

                // hdlr
                write_full_box(buf, b"hdlr", 0, 0, |buf| {
                    buf.put_u32(0); // pre_defined
                    buf.put_slice(b"vide"); // handler_type
                    buf.put_bytes(0, 12); // reserved
                    buf.put_slice(b"LVQR Video\0");
                });

                // minf
                write_box(buf, b"minf", |buf| {
                    // vmhd
                    write_full_box(buf, b"vmhd", 0, 1, |buf| {
                        buf.put_u16(0); // graphicsmode
                        buf.put_bytes(0, 6); // opcolor
                    });

                    // dinf
                    write_box(buf, b"dinf", |buf| {
                        write_full_box(buf, b"dref", 0, 0, |buf| {
                            buf.put_u32(1); // entry_count
                            write_full_box(buf, b"url ", 0, 1, |_buf| {
                                // self-contained flag set, no URL
                            });
                        });
                    });

                    // stbl
                    write_box(buf, b"stbl", |buf| {
                        // stsd
                        write_full_box(buf, b"stsd", 0, 0, |buf| {
                            buf.put_u32(1); // entry_count

                            // avc1 sample entry
                            write_box(buf, b"avc1", |buf| {
                                buf.put_bytes(0, 6); // reserved
                                buf.put_u16(1); // data_reference_index
                                buf.put_bytes(0, 16); // pre_defined + reserved
                                buf.put_u16(width);
                                buf.put_u16(height);
                                buf.put_u32(0x00480000); // horizresolution = 72 dpi
                                buf.put_u32(0x00480000); // vertresolution = 72 dpi
                                buf.put_u32(0); // reserved
                                buf.put_u16(1); // frame_count
                                buf.put_bytes(0, 32); // compressorname
                                buf.put_u16(0x0018); // depth = 24
                                buf.put_i16(-1); // pre_defined

                                // avcC box (AVCDecoderConfigurationRecord)
                                write_box(buf, b"avcC", |buf| {
                                    buf.put_u8(1); // configurationVersion
                                    buf.put_u8(config.profile);
                                    buf.put_u8(config.compat);
                                    buf.put_u8(config.level);
                                    buf.put_u8(0xFF); // lengthSizeMinusOne=3 | reserved
                                    buf.put_u8(0xE0 | (config.sps_list.len() as u8)); // numSPS | reserved
                                    for sps in &config.sps_list {
                                        buf.put_u16(sps.len() as u16);
                                        buf.put_slice(sps);
                                    }
                                    buf.put_u8(config.pps_list.len() as u8); // numPPS
                                    for pps in &config.pps_list {
                                        buf.put_u16(pps.len() as u16);
                                        buf.put_slice(pps);
                                    }
                                });
                            });
                        });

                        // Empty required boxes
                        write_full_box(buf, b"stts", 0, 0, |buf| buf.put_u32(0));
                        write_full_box(buf, b"stsc", 0, 0, |buf| buf.put_u32(0));
                        write_full_box(buf, b"stsz", 0, 0, |buf| {
                            buf.put_u32(0); // sample_size
                            buf.put_u32(0); // sample_count
                        });
                        write_full_box(buf, b"stco", 0, 0, |buf| buf.put_u32(0));
                    });
                });
            });
        });

        // mvex
        write_box(buf, b"mvex", |buf| {
            write_full_box(buf, b"trex", 0, 0, |buf| {
                buf.put_u32(1); // track_ID
                buf.put_u32(1); // default_sample_description_index
                buf.put_u32(0); // default_sample_duration
                buf.put_u32(0); // default_sample_size
                buf.put_u32(0); // default_sample_flags
            });
        });
    });

    buf.freeze()
}

/// Generate an audio init segment (ftyp + moov) for AAC.
pub fn audio_init_segment(config: &AudioConfig) -> Bytes {
    let mut buf = BytesMut::with_capacity(512);

    // ftyp
    write_box(&mut buf, b"ftyp", |buf| {
        buf.put_slice(b"isom");
        buf.put_u32(0);
        buf.put_slice(b"isom");
        buf.put_slice(b"iso6");
        buf.put_slice(b"msdh");
        buf.put_slice(b"msix");
    });

    // moov
    write_box(&mut buf, b"moov", |buf| {
        // mvhd
        write_full_box(buf, b"mvhd", 0, 0, |buf| {
            buf.put_u32(0);
            buf.put_u32(0);
            buf.put_u32(config.sample_rate);
            buf.put_u32(0);
            buf.put_u32(0x00010000);
            buf.put_u16(0x0100);
            buf.put_bytes(0, 10);
            for &v in &[0x00010000u32, 0, 0, 0, 0x00010000, 0, 0, 0, 0x40000000] {
                buf.put_u32(v);
            }
            buf.put_bytes(0, 24);
            buf.put_u32(2);
        });

        // trak
        write_box(buf, b"trak", |buf| {
            write_full_box(buf, b"tkhd", 0, 0x03, |buf| {
                buf.put_u32(0);
                buf.put_u32(0);
                buf.put_u32(1);
                buf.put_u32(0);
                buf.put_u32(0);
                buf.put_bytes(0, 8);
                buf.put_u16(0);
                buf.put_u16(0);
                buf.put_u16(0x0100); // volume = 1.0 (audio)
                buf.put_u16(0);
                for &v in &[0x00010000u32, 0, 0, 0, 0x00010000, 0, 0, 0, 0x40000000] {
                    buf.put_u32(v);
                }
                buf.put_u32(0);
                buf.put_u32(0);
            });

            write_box(buf, b"mdia", |buf| {
                write_full_box(buf, b"mdhd", 0, 0, |buf| {
                    buf.put_u32(0);
                    buf.put_u32(0);
                    buf.put_u32(config.sample_rate);
                    buf.put_u32(0);
                    buf.put_u32(0x55C40000);
                });

                write_full_box(buf, b"hdlr", 0, 0, |buf| {
                    buf.put_u32(0);
                    buf.put_slice(b"soun");
                    buf.put_bytes(0, 12);
                    buf.put_slice(b"LVQR Audio\0");
                });

                write_box(buf, b"minf", |buf| {
                    write_full_box(buf, b"smhd", 0, 0, |buf| {
                        buf.put_u16(0); // balance
                        buf.put_u16(0); // reserved
                    });

                    write_box(buf, b"dinf", |buf| {
                        write_full_box(buf, b"dref", 0, 0, |buf| {
                            buf.put_u32(1);
                            write_full_box(buf, b"url ", 0, 1, |_buf| {});
                        });
                    });

                    write_box(buf, b"stbl", |buf| {
                        write_full_box(buf, b"stsd", 0, 0, |buf| {
                            buf.put_u32(1);

                            // mp4a sample entry
                            write_box(buf, b"mp4a", |buf| {
                                buf.put_bytes(0, 6); // reserved
                                buf.put_u16(1); // data_reference_index
                                buf.put_bytes(0, 8); // reserved
                                buf.put_u16(config.channels as u16);
                                buf.put_u16(16); // sampleSize
                                buf.put_u16(0); // pre_defined
                                buf.put_u16(0); // reserved
                                buf.put_u32(config.sample_rate << 16); // sampleRate (fixed-point 16.16)

                                // esds box (ISO/IEC 14496-1 §8). The ASC bytes in
                                // `config.asc` are the DecoderSpecificInfo payload; they
                                // were validated up front by
                                // `lvqr_codec::aac::parse_asc` when the FLV AAC sequence
                                // header arrived, so any bit layout supported by the
                                // hardened parser (xHE-AAC, HE-AAC SBR/PS, explicit-
                                // frequency escape, >127-byte ASC) now round-trips
                                // through the writer without the previous single-byte
                                // descriptor-length footgun.
                                write_full_box(buf, b"esds", 0, 0, |buf| {
                                    write_mpeg4_descriptor(buf, 0x03, |buf| {
                                        // ES_ID + flags
                                        buf.put_u16(1);
                                        buf.put_u8(0);

                                        // DecoderConfigDescriptor
                                        write_mpeg4_descriptor(buf, 0x04, |buf| {
                                            buf.put_u8(0x40); // objectTypeIndication = Audio ISO/IEC 14496-3
                                            buf.put_u8(0x15); // streamType=5 (audio) | upStream=0 | reserved=1
                                            buf.put_u8(0); // bufferSizeDB (24 bits)
                                            buf.put_u16(0);
                                            buf.put_u32(0); // maxBitrate
                                            buf.put_u32(0); // avgBitrate

                                            // DecoderSpecificInfo wraps the ASC bytes.
                                            write_mpeg4_descriptor(buf, 0x05, |buf| {
                                                buf.put_slice(&config.asc);
                                            });
                                        });

                                        // SLConfigDescriptor
                                        write_mpeg4_descriptor(buf, 0x06, |buf| {
                                            buf.put_u8(0x02); // predefined = MP4
                                        });
                                    });
                                });
                            });
                        });

                        write_full_box(buf, b"stts", 0, 0, |buf| buf.put_u32(0));
                        write_full_box(buf, b"stsc", 0, 0, |buf| buf.put_u32(0));
                        write_full_box(buf, b"stsz", 0, 0, |buf| {
                            buf.put_u32(0);
                            buf.put_u32(0);
                        });
                        write_full_box(buf, b"stco", 0, 0, |buf| buf.put_u32(0));
                    });
                });
            });
        });

        write_box(buf, b"mvex", |buf| {
            write_full_box(buf, b"trex", 0, 0, |buf| {
                buf.put_u32(1);
                buf.put_u32(1);
                buf.put_u32(0);
                buf.put_u32(0);
                buf.put_u32(0);
            });
        });
    });

    buf.freeze()
}

// --- Media segment generation ---
//
// The hand-rolled video media-segment writer lived here through v0.3.
// Session 14 removed it in favour of routing every per-frame RTMP video
// payload straight through `lvqr_cmaf::build_moof_mdat`; the AVC parity
// gate had kept the two paths aligned across sessions 7-13, and the
// cmaf-writer default-on flip in session 12.2 completed the transition.
// The audio media-segment writer below is unrelated and stays.

/// Generate an audio media segment (moof + mdat) containing a single AAC frame.
pub fn audio_segment(sequence: u32, base_dts: u64, duration: u32, data: &Bytes) -> Bytes {
    let mut buf = BytesMut::with_capacity(128 + data.len());

    let moof_start = buf.len();
    write_box(&mut buf, b"moof", |buf| {
        write_full_box(buf, b"mfhd", 0, 0, |buf| {
            buf.put_u32(sequence);
        });

        write_box(buf, b"traf", |buf| {
            write_full_box(buf, b"tfhd", 0, 0x020000, |buf| {
                buf.put_u32(1);
            });

            write_full_box(buf, b"tfdt", 1, 0, |buf| {
                buf.put_u64(base_dts);
            });

            // trun: data_offset + duration + size
            let trun_flags: u32 = 0x000001 | 0x000100 | 0x000200;
            write_full_box(buf, b"trun", 0, trun_flags, |buf| {
                buf.put_u32(1); // sample_count
                buf.put_i32(0); // data_offset placeholder
                buf.put_u32(duration);
                buf.put_u32(data.len() as u32);
            });
        });
    });

    let moof_size = buf.len() - moof_start;
    let data_offset = (moof_size + 8) as i32;
    patch_trun_data_offset(&mut buf, moof_start, data_offset);

    write_box(&mut buf, b"mdat", |buf| {
        buf.put_slice(data);
    });

    buf.freeze()
}

/// Find the trun box within a moof and patch its data_offset field.
fn patch_trun_data_offset(buf: &mut BytesMut, moof_start: usize, data_offset: i32) {
    let mut pos = moof_start + 8; // skip moof header
    while pos + 8 <= buf.len() {
        let box_size = u32::from_be_bytes([buf[pos], buf[pos + 1], buf[pos + 2], buf[pos + 3]]) as usize;
        let box_type = &buf[pos + 4..pos + 8];

        if box_type == b"traf" {
            // Recurse into traf
            let traf_end = pos + box_size;
            let mut inner = pos + 8;
            while inner + 8 <= traf_end {
                let inner_size =
                    u32::from_be_bytes([buf[inner], buf[inner + 1], buf[inner + 2], buf[inner + 3]]) as usize;
                let inner_type = &buf[inner + 4..inner + 8];

                if inner_type == b"trun" {
                    // trun: [4 size][4 type][4 full_box_header][4 sample_count][4 data_offset]
                    let offset_pos = inner + 8 + 4 + 4; // after header + full_box + sample_count
                    let bytes = data_offset.to_be_bytes();
                    buf[offset_pos..offset_pos + 4].copy_from_slice(&bytes);
                    return;
                }
                inner += inner_size;
            }
        }
        pos += box_size;
    }
}

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

    fn read_box_at(data: &[u8], offset: usize) -> Option<(usize, &[u8; 4], &[u8])> {
        if offset + 8 > data.len() {
            return None;
        }
        let size = u32::from_be_bytes([data[offset], data[offset + 1], data[offset + 2], data[offset + 3]]) as usize;
        if offset + size > data.len() || size < 8 {
            return None;
        }
        let box_type: &[u8; 4] = data[offset + 4..offset + 8].try_into().ok()?;
        let payload = &data[offset + 8..offset + size];
        Some((size, box_type, payload))
    }

    fn find_box<'a>(data: &'a [u8], target: &[u8; 4]) -> Option<(usize, &'a [u8])> {
        let mut pos = 0;
        while let Some((size, box_type, payload)) = read_box_at(data, pos) {
            if box_type == target {
                return Some((pos, payload));
            }
            pos += size;
        }
        None
    }

    fn test_video_config() -> flv::VideoConfig {
        flv::VideoConfig {
            sps_list: vec![vec![0x67, 0x64, 0x00, 0x1F, 0xAC, 0xD9]],
            pps_list: vec![vec![0x68, 0xEE, 0x3C, 0x80]],
            profile: 0x64,
            compat: 0x00,
            level: 0x1F,
            nalu_length_size: 4,
        }
    }

    fn test_audio_config() -> flv::AudioConfig {
        flv::AudioConfig {
            asc: vec![0x12, 0x10], // AAC-LC, 44100 Hz, stereo
            sample_rate: 44100,
            channels: 2,
            object_type: 2,
        }
    }

    #[test]
    fn video_init_starts_with_ftyp() {
        let init = video_init_segment(&test_video_config());
        assert!(init.len() > 8);
        assert_eq!(&init[4..8], b"ftyp");
    }

    #[test]
    fn video_init_contains_moov() {
        let init = video_init_segment(&test_video_config());
        assert!(find_box(&init, b"moov").is_some());
    }

    #[test]
    fn video_init_contains_avcc_with_sps_pps() {
        let config = test_video_config();
        let init = video_init_segment(&config);

        // Search for avcC box by scanning the bytes
        let init_bytes = &init[..];
        let avcc_needle = b"avcC";
        let pos = init_bytes
            .windows(4)
            .position(|w| w == avcc_needle)
            .expect("avcC box not found");

        // avcC payload starts after the box header (4 bytes before the type + 4 type bytes = skip 4 after type)
        let avcc_start = pos + 4; // past "avcC" type
        assert_eq!(init_bytes[avcc_start], 1); // configurationVersion
        assert_eq!(init_bytes[avcc_start + 1], config.profile);
        assert_eq!(init_bytes[avcc_start + 2], config.compat);
        assert_eq!(init_bytes[avcc_start + 3], config.level);
    }

    #[test]
    fn audio_init_starts_with_ftyp() {
        let init = audio_init_segment(&test_audio_config());
        assert_eq!(&init[4..8], b"ftyp");
    }

    #[test]
    fn audio_init_contains_esds() {
        let init = audio_init_segment(&test_audio_config());
        let init_bytes = &init[..];
        let esds_pos = init_bytes.windows(4).position(|w| w == b"esds");
        assert!(esds_pos.is_some(), "esds box not found in audio init");
    }

    #[test]
    fn mpeg4_descriptor_length_encoding_round_trips_large_payloads() {
        // The esds writer migration swapped a single-byte descriptor
        // length prefix for the full 4-byte MPEG-4 variable-length
        // form. The unit tests above only exercise the 2-byte ASC
        // case (payload well under 127 bytes). This test verifies
        // that payloads >= 128 bytes round-trip through the length
        // encoding without corruption, which is the whole reason the
        // migration exists.
        //
        // The payload is 200 bytes of a non-zero pattern so every
        // byte position is distinguishable from the padding a buggy
        // length field might introduce.
        let payload: Vec<u8> = (0..200u8).collect();
        let mut buf = BytesMut::new();
        write_mpeg4_descriptor(&mut buf, 0x05, |inner| {
            inner.put_slice(&payload);
        });

        // Tag byte + 4 length bytes + 200 payload bytes = 205 total.
        assert_eq!(buf.len(), 1 + 4 + payload.len());
        assert_eq!(buf[0], 0x05, "tag byte");
        // Expected encoding for size 200 = 0xC8:
        //   byte 0: 0x80 | ((200 >> 21) & 0x7F) = 0x80
        //   byte 1: 0x80 | ((200 >> 14) & 0x7F) = 0x80
        //   byte 2: 0x80 | ((200 >>  7) & 0x7F) = 0x81
        //   byte 3:         200 & 0x7F          = 0x48
        assert_eq!(buf[1], 0x80);
        assert_eq!(buf[2], 0x80);
        assert_eq!(buf[3], 0x81);
        assert_eq!(buf[4], 0x48);
        assert_eq!(&buf[5..], payload.as_slice());
    }

    #[test]
    fn audio_segment_structure() {
        let data = Bytes::from(vec![0x01, 0x02, 0x03, 0x04]);
        let seg = audio_segment(1, 0, 1024, &data);

        assert_eq!(&seg[4..8], b"moof");
        let moof_size = u32::from_be_bytes([seg[0], seg[1], seg[2], seg[3]]) as usize;
        assert_eq!(&seg[moof_size + 4..moof_size + 8], b"mdat");
    }
}