path: root/remuxer/mp4/src/boxes.rs

/*
    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
    Copyright (C) 2026 metamuffin <metamuffin.org>
*/

use crate::{BoxW, WriteBox};

macro_rules! container_box {
    ($typ:tt, $iden:literal) => {
        pub struct $typ<T: Fn(&mut BoxW)>(pub T);
        impl<T: Fn(&mut BoxW)> WriteBox for $typ<T> {
            const BOXTYPE: [u8; 4] = *$iden;
            fn write(&self, buf: &mut Vec<u8>) {
                self.0(&mut BoxW(buf))
            }
        }
    };
}

container_box!(Movie, b"moov");
container_box!(Track, b"trak");
container_box!(Edit, b"edts");
container_box!(Media, b"mdia");
container_box!(MediaInformation, b"minf");
container_box!(DataInformation, b"dinf");
container_box!(SampleTable, b"stbl");
container_box!(MovieFragment, b"moof");
container_box!(TrackFragment, b"traf");
container_box!(MovieExtends, b"mvex");

pub struct FileType<'a> {
    pub major_brand: [u8; 4],
    pub minor_version: u32,
    pub compatible_brands: &'a [[u8; 4]],
}
impl WriteBox for FileType<'_> {
    const BOXTYPE: [u8; 4] = *b"ftyp";
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.major_brand);
        buf.extend(self.minor_version.to_be_bytes());
        for cb in self.compatible_brands {
            buf.extend(cb);
        }
    }
}

pub struct SegmentType<'a> {
    pub major_brand: [u8; 4],
    pub minor_version: u32,
    pub compatible_brands: &'a [[u8; 4]],
}
impl WriteBox for SegmentType<'_> {
    const BOXTYPE: [u8; 4] = *b"styp";
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.major_brand);
        buf.extend(self.minor_version.to_be_bytes());
        for cb in self.compatible_brands {
            buf.extend(cb);
        }
    }
}

pub struct MovieHeader {
    pub creation_time: u64,
    pub modification_time: u64,
    pub timescale: u32,
    pub duration: u64,
    pub rate: u32,
    pub volume: u16,
    pub matrix: [u32; 9],
    pub pre_defined: [u32; 6],
    pub next_track_id: u32,
}

impl WriteBox for MovieHeader {
    const BOXTYPE: [u8; 4] = *b"mvhd";
    const VERSION: Option<u8> = Some(1);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.creation_time.to_be_bytes());
        buf.extend(self.modification_time.to_be_bytes());
        buf.extend(self.timescale.to_be_bytes());
        buf.extend(self.duration.to_be_bytes());
        buf.extend(self.rate.to_be_bytes());
        buf.extend(self.volume.to_be_bytes());
        buf.extend([0; 10]);
        for e in self.matrix {
            buf.extend(e.to_be_bytes());
        }
        buf.extend([0; 24]);
        buf.extend(self.next_track_id.to_be_bytes());
    }
}
impl Default for MovieHeader {
    fn default() -> Self {
        Self {
            creation_time: 0,
            modification_time: 0,
            timescale: 0,
            duration: 0,
            rate: 0x00010000,
            volume: 0x0100,
            matrix: [0x00010000, 0, 0, 0, 0x00010000, 0, 0, 0, 0x00010000],
            pre_defined: [0; 6],
            next_track_id: 0,
        }
    }
}

pub struct TrackHeader {
    pub creation_time: u64,
    pub modification_time: u64,
    pub track_id: u32,
    pub duration: u64,
    pub layer: u16,
    pub alternate_group: u16,
    pub volume: u16,
    pub matrix: [u32; 9],
    pub width: u32,
    pub height: u32,
}
impl WriteBox for TrackHeader {
    const BOXTYPE: [u8; 4] = *b"tkhd";
    const VERSION: Option<u8> = Some(1);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.creation_time.to_be_bytes());
        buf.extend(self.modification_time.to_be_bytes());
        buf.extend(self.track_id.to_be_bytes());
        buf.extend([0; 4]);
        buf.extend(self.duration.to_be_bytes());
        buf.extend([0; 8]);
        buf.extend(self.layer.to_be_bytes());
        buf.extend(self.alternate_group.to_be_bytes());
        buf.extend(self.volume.to_be_bytes());
        buf.extend([0; 2]);
        for e in self.matrix {
            buf.extend(e.to_be_bytes());
        }
        buf.extend(self.width.to_be_bytes());
        buf.extend(self.height.to_be_bytes());
    }
}
impl Default for TrackHeader {
    fn default() -> Self {
        Self {
            creation_time: 0,
            modification_time: 0,
            track_id: 0,
            duration: 0,
            layer: 0,
            alternate_group: 0,
            volume: 0,
            matrix: [0x00010000, 0, 0, 0, 0x00010000, 0, 0, 0, 0x00010000],
            width: 0,
            height: 0,
        }
    }
}

pub struct EditList<'a>(pub &'a [EditListEntry]);
pub struct EditListEntry {
    pub edit_duration: u64,
    pub media_time: u64,
    pub media_rate: u32,
}
impl WriteBox for EditList<'_> {
    const BOXTYPE: [u8; 4] = *b"elst";
    const VERSION: Option<u8> = Some(1);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend((self.0.len() as u32).to_be_bytes());
        for e in self.0 {
            buf.extend(e.edit_duration.to_be_bytes());
            buf.extend(e.media_time.to_be_bytes());
            buf.extend(e.media_rate.to_be_bytes());
        }
    }
}

#[derive(Default)]
pub struct MediaHeader {
    pub creation_time: u64,
    pub modification_time: u64,
    pub timescale: u32,
    pub duration: u64,
    pub language: [u8; 3],
    pub pre_defined: u16,
}
impl WriteBox for MediaHeader {
    const BOXTYPE: [u8; 4] = *b"mdhd";
    const VERSION: Option<u8> = Some(1);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.creation_time.to_be_bytes());
        buf.extend(self.modification_time.to_be_bytes());
        buf.extend(self.timescale.to_be_bytes());
        buf.extend(self.duration.to_be_bytes());
        buf.extend(
            ((self.language[0] as u16 & 0b11111 << 10)
                | (self.language[1] as u16 & 0b11111 << 5)
                | (self.language[2] as u16 & 0b11111))
                .to_be_bytes(),
        );
        buf.extend([0; 2]);
    }
}

pub struct Handler<'a> {
    pub handler_type: [u8; 4],
    pub name: &'a str,
}
impl WriteBox for Handler<'_> {
    const BOXTYPE: [u8; 4] = *b"hdlr";
    const VERSION: Option<u8> = Some(0);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend([0; 4]);
        buf.extend(self.handler_type);
        buf.extend([0; 12]);
        buf.extend(self.name.as_bytes());
        buf.push(0);
    }
}

#[derive(Default)]
pub struct VideoMediaHeader {
    pub graphicsmode: u16,
    pub opcolor: [u16; 3],
}
impl WriteBox for VideoMediaHeader {
    const BOXTYPE: [u8; 4] = *b"vmhd";
    const VERSION: Option<u8> = Some(0);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.graphicsmode.to_be_bytes());
        buf.extend(self.opcolor[0].to_be_bytes());
        buf.extend(self.opcolor[1].to_be_bytes());
        buf.extend(self.opcolor[2].to_be_bytes());
    }
}

#[derive(Default)]
pub struct SoundMediaHeader {
    pub balance: u16,
}
impl WriteBox for SoundMediaHeader {
    const BOXTYPE: [u8; 4] = *b"smhd";
    const VERSION: Option<u8> = Some(0);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.balance.to_be_bytes());
    }
}

pub struct DataReference<'a>(pub &'a [DataReferenceEntry]);
pub enum DataReferenceEntry {
    URL(String),
    URN { name: String, location: String },
    Imda(u32),
    SeqNumImda,
}
impl WriteBox for DataReference<'_> {
    const BOXTYPE: [u8; 4] = *b"dref";
    const VERSION: Option<u8> = Some(0);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend((self.0.len() as u32).to_be_bytes());
        for e in self.0 {
            match e {
                DataReferenceEntry::URL(x) if x.is_empty() => {
                    buf.extend(12u32.to_be_bytes());
                    buf.extend(b"url ");
                    buf.push(0); // version
                    buf.extend([0, 0, 1]); // flags
                }
                _ => todo!(),
            }
        }
    }
}

pub struct MovieFragmentHeader {
    pub sequence_number: u32,
}
impl WriteBox for MovieFragmentHeader {
    const BOXTYPE: [u8; 4] = *b"mfhd";
    const VERSION: Option<u8> = Some(0);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.sequence_number.to_be_bytes());
    }
}

pub struct TrackFragmentHeader {
    pub track_id: u32,
}
impl WriteBox for TrackFragmentHeader {
    const BOXTYPE: [u8; 4] = *b"tfhd";
    const VERSION: Option<u8> = Some(0);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.track_id.to_be_bytes());
    }
}

pub struct TrackFragmentBaseMediaDecodeTime {
    pub base_media_decode_time: u64,
}
impl WriteBox for TrackFragmentBaseMediaDecodeTime {
    const BOXTYPE: [u8; 4] = *b"tfdt";
    const VERSION: Option<u8> = Some(1);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.base_media_decode_time.to_be_bytes());
    }
}

pub struct TrackRun<'a> {
    pub data_offset: i32,
    pub samples: &'a [TrackRunSample],
}
pub struct TrackRunSample {
    pub duration: u32,
    pub size: u32,
    pub flags: u32,
    pub composition_time_offset: i32,
}
impl WriteBox for TrackRun<'_> {
    const BOXTYPE: [u8; 4] = *b"trun";
    const VERSION: Option<u8> = Some(1);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend((self.samples.len() as u32).to_be_bytes());
        buf.extend(self.data_offset.to_be_bytes());
        for s in self.samples {
            buf.extend(s.duration.to_be_bytes());
            buf.extend(s.size.to_be_bytes());
            buf.extend(s.flags.to_be_bytes());
            buf.extend(s.composition_time_offset.to_be_bytes());
        }
    }
    fn flags(&self) -> u32 {
        0x000001 // data-offset-present
        | 0x000100 // sample-duration-present
        | 0x000200 // sample-size-present
        | 0x000400 // sample-flags-present
        | 0x000800 // sample-composition-time-offsets-present
    }
}

pub struct TimeToSample<'a>(pub &'a [TimeToSampleEntry]);
pub struct TimeToSampleEntry {
    pub sample_count: u32,
    pub sample_delta: u32,
}
impl WriteBox for TimeToSample<'_> {
    const BOXTYPE: [u8; 4] = *b"stts";
    const VERSION: Option<u8> = Some(0);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend((self.0.len() as u32).to_be_bytes());
        for e in self.0 {
            buf.extend(e.sample_count.to_be_bytes());
            buf.extend(e.sample_delta.to_be_bytes());
        }
    }
}

pub struct TimeToChunk<'a>(pub &'a [TimeToChunkEntry]);
pub struct TimeToChunkEntry {
    pub first_chunk: u32,
    pub samples_per_chunk: u32,
    pub sample_description_index: u32,
}
impl WriteBox for TimeToChunk<'_> {
    const BOXTYPE: [u8; 4] = *b"stsc";
    const VERSION: Option<u8> = Some(0);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend((self.0.len() as u32).to_be_bytes());
        for e in self.0 {
            buf.extend(e.first_chunk.to_be_bytes());
            buf.extend(e.samples_per_chunk.to_be_bytes());
            buf.extend(e.sample_description_index.to_be_bytes());
        }
    }
}

pub struct SampleSize<'a> {
    pub sample_size: u32,
    pub entries: &'a [u32],
}
impl WriteBox for SampleSize<'_> {
    const BOXTYPE: [u8; 4] = *b"stts";
    const VERSION: Option<u8> = Some(0);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.sample_size.to_be_bytes());
        buf.extend((self.entries.len() as u32).to_be_bytes());
        if self.sample_size == 0 {
            for e in self.entries {
                buf.extend(e.to_be_bytes());
            }
        }
    }
}

pub struct ChunkOffset<'a>(pub &'a [u32]);
impl WriteBox for ChunkOffset<'_> {
    const BOXTYPE: [u8; 4] = *b"stco";
    const VERSION: Option<u8> = Some(0);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend((self.0.len() as u32).to_be_bytes());
        for e in self.0 {
            buf.extend(e.to_be_bytes());
        }
    }
}

pub struct SyncSample<'a>(pub &'a [u32]);
impl WriteBox for SyncSample<'_> {
    const BOXTYPE: [u8; 4] = *b"stss";
    const VERSION: Option<u8> = Some(0);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend((self.0.len() as u32).to_be_bytes());
        for e in self.0 {
            buf.extend(e.to_be_bytes());
        }
    }
}

pub struct SampleDescription<T: Fn(&mut BoxW)>(pub T);
impl<T: Fn(&mut BoxW)> WriteBox for SampleDescription<T> {
    const BOXTYPE: [u8; 4] = *b"stsd";
    const VERSION: Option<u8> = Some(0);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(1u32.to_be_bytes());
        self.0(&mut BoxW(buf))
    }
}

pub struct AVCSampleEntry<'a> {
    pub visual: VisualSampleEntry,
    pub config: AVCConfiguration<'a>,
}
impl WriteBox for AVCSampleEntry<'_> {
    const BOXTYPE: [u8; 4] = *b"avc1";
    fn write(&self, buf: &mut Vec<u8>) {
        self.visual.write(buf);
        BoxW(buf).write(&self.config);
        if let Some(colr) = &self.visual.colr {
            BoxW(buf).write(colr);
        }
        if let Some(pasp) = &self.visual.pasp {
            BoxW(buf).write(pasp);
        }
    }
}

pub struct AVCConfiguration<'a>(pub &'a [u8]);
impl WriteBox for AVCConfiguration<'_> {
    const BOXTYPE: [u8; 4] = *b"avcC";
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.0);
    }
}

#[derive(Default)]
pub struct SampleEntry {
    pub data_reference_index: u16,
}
impl WriteBox for SampleEntry {
    const BOXTYPE: [u8; 4] = [0; 4];
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend([0; 6]);
        buf.extend(self.data_reference_index.to_be_bytes())
    }
}

pub struct VisualSampleEntry {
    pub sample: SampleEntry,
    pub width: u16,
    pub height: u16,
    pub horizresolution: u32,
    pub vertresolution: u32,
    pub frame_count: u16,
    pub compressorname: [u8; 32],
    pub depth: u16,
    pub pasp: Option<PixelAspectRatio>,
    pub colr: Option<ColourInformation>,
}
impl WriteBox for VisualSampleEntry {
    const BOXTYPE: [u8; 4] = [0; 4];
    fn write(&self, buf: &mut Vec<u8>) {
        self.sample.write(buf);
        buf.extend([0; 2]);
        buf.extend([0; 2]);
        buf.extend([0; 12]);
        buf.extend(self.width.to_be_bytes());
        buf.extend(self.height.to_be_bytes());
        buf.extend(self.horizresolution.to_be_bytes());
        buf.extend(self.vertresolution.to_be_bytes());
        buf.extend([0; 4]);
        buf.extend(self.frame_count.to_be_bytes());
        buf.extend(self.compressorname);
        buf.extend(self.depth.to_be_bytes());
        buf.extend((-1i16).to_be_bytes());
    }
}
impl Default for VisualSampleEntry {
    fn default() -> Self {
        Self {
            sample: SampleEntry::default(),
            width: 0,
            height: 0,
            horizresolution: 0x00480000,
            vertresolution: 0x00480000,
            frame_count: 1,
            compressorname: [0; 32],
            depth: 0x0018,
            pasp: None,
            colr: None,
        }
    }
}

pub struct PixelAspectRatio {
    pub h_spacing: u32,
    pub v_spacing: u32,
}
impl WriteBox for PixelAspectRatio {
    const BOXTYPE: [u8; 4] = *b"pasp";
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.h_spacing.to_be_bytes());
        buf.extend(self.v_spacing.to_be_bytes());
    }
}

pub enum ColourInformation {
    OnScreenColours {
        colour_primaries: u16,
        transfer_charateristics: u16,
        matrix_coefficients: u16,
        full_range: bool,
    },
    RestrictedICC,
    UnrestrictedICC,
}
impl WriteBox for ColourInformation {
    const BOXTYPE: [u8; 4] = *b"colr";
    fn write(&self, buf: &mut Vec<u8>) {
        match self {
            ColourInformation::OnScreenColours {
                colour_primaries,
                transfer_charateristics,
                matrix_coefficients,
                full_range,
            } => {
                buf.extend(b"nclx");
                buf.extend(colour_primaries.to_be_bytes());
                buf.extend(transfer_charateristics.to_be_bytes());
                buf.extend(matrix_coefficients.to_be_bytes());
                buf.push((*full_range as u8) << 7);
            }
            ColourInformation::RestrictedICC => {
                buf.extend(b"rICC");
                todo!()
            }
            ColourInformation::UnrestrictedICC => {
                buf.extend(b"prof");
                todo!()
            }
        }
    }
}

pub struct MovieExtendsHeader {
    pub fragment_duration: u64,
}
impl WriteBox for MovieExtendsHeader {
    const BOXTYPE: [u8; 4] = *b"mehd";
    const VERSION: Option<u8> = Some(1);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.fragment_duration.to_be_bytes())
    }
}

#[derive(Default)]
pub struct TrackExtends {
    pub track_id: u32,
    pub default_sample_description_index: u32,
    pub default_sample_duration: u32,
    pub default_sample_size: u32,
    pub default_sample_flags: u32,
}
impl WriteBox for TrackExtends {
    const BOXTYPE: [u8; 4] = *b"trex";
    const VERSION: Option<u8> = Some(0);
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.track_id.to_be_bytes());
        buf.extend(self.default_sample_description_index.to_be_bytes());
        buf.extend(self.default_sample_duration.to_be_bytes());
        buf.extend(self.default_sample_size.to_be_bytes());
        buf.extend(self.default_sample_flags.to_be_bytes());
    }
}

pub struct MediaData<'a>(pub &'a [u8]);
impl WriteBox for MediaData<'_> {
    const BOXTYPE: [u8; 4] = *b"mdat";
    fn write(&self, buf: &mut Vec<u8>) {
        buf.extend(self.0);
    }
}