several mp4 fixes

7 files changed, 123 insertions, 21 deletions

diff --git a/remuxer/mp4/src/boxes.rs b/remuxer/mp4/src/boxes.rs
index 5a4e108..3a1b188 100644
--- a/remuxer/mp4/src/boxes.rs
+++ b/remuxer/mp4/src/boxes.rs
@@ -25,8 +25,8 @@ container_box!(Media, b"mdia");
 container_box!(MediaInformation, b"minf");
 container_box!(DataInformation, b"dinf");
 container_box!(SampleTable, b"stbl");
-container_box!(MovieFragment, b"stbl");
-container_box!(TrackFragment, b"stbl");
+container_box!(MovieFragment, b"moof");
+container_box!(TrackFragment, b"traf");
 container_box!(MovieExtends, b"mvex");
 
 pub struct FileType<'a> {
@@ -306,14 +306,35 @@ impl WriteBox for TrackFragmentBaseMediaDecodeTime {
     }
 }
 
-pub struct TrackRun {
+pub struct TrackRun<'a> {
     pub data_offset: i32,
+    pub samples: &'a [TrackRunSample],
 }
-impl WriteBox for TrackRun {
+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
     }
 }
 
@@ -570,3 +591,11 @@ impl WriteBox for TrackExtends {
         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);
+    }
+}
diff --git a/remuxer/mp4/src/lib.rs b/remuxer/mp4/src/lib.rs
index 7af8417..2098e1b 100644
--- a/remuxer/mp4/src/lib.rs
+++ b/remuxer/mp4/src/lib.rs
@@ -12,6 +12,9 @@ impl<'a> BoxW<'a> {
     pub fn new(buffer: &'a mut Vec<u8>) -> Self {
         Self(buffer)
     }
+    pub fn buffer_mut(&mut self) -> &mut Vec<u8> {
+        self.0
+    }
     pub fn write<T: WriteBox>(&mut self, b: T) {
         let start = self.0.len();
         self.0.extend(0u32.to_be_bytes());
diff --git a/remuxer/src/muxers/matroska.rs b/remuxer/src/muxers/matroska.rs
index 5a2d024..272c91d 100644
--- a/remuxer/src/muxers/matroska.rs
+++ b/remuxer/src/muxers/matroska.rs
@@ -68,7 +68,7 @@ impl FragmentMuxer for MatroskaFragmentMuxer {
     fn write_init(out: &mut Vec<u8>, segment: Segment) -> Result<()> {
         write_init_shared(out, segment, false)
     }
-    fn write_frag(out: &mut Vec<u8>, segment: Segment) -> Result<()> {
+    fn write_frag(out: &mut Vec<u8>, segment: Segment, _index: usize, _end_pts: u64) -> Result<()> {
         write_frag_shared(out, segment, false)
     }
 }
@@ -77,7 +77,7 @@ impl FragmentMuxer for WebmFragmentMuxer {
     fn write_init(out: &mut Vec<u8>, segment: Segment) -> Result<()> {
         write_init_shared(out, segment, true)
     }
-    fn write_frag(out: &mut Vec<u8>, segment: Segment) -> Result<()> {
+    fn write_frag(out: &mut Vec<u8>, segment: Segment, _index: usize, _end_pts: u64) -> Result<()> {
         write_frag_shared(out, segment, true)
     }
 }
diff --git a/remuxer/src/muxers/mod.rs b/remuxer/src/muxers/mod.rs
index d668ab2..3d0a935 100644
--- a/remuxer/src/muxers/mod.rs
+++ b/remuxer/src/muxers/mod.rs
@@ -19,7 +19,7 @@ use winter_matroska::Segment;
 
 pub trait FragmentMuxer {
     fn write_init(out: &mut Vec<u8>, segment: Segment) -> Result<()>;
-    fn write_frag(out: &mut Vec<u8>, segment: Segment) -> Result<()>;
+    fn write_frag(out: &mut Vec<u8>, segment: Segment, index: usize, end_pts: u64) -> Result<()>;
 }
 
 pub fn write_init(container: ContainerFormat, out: &mut Vec<u8>, segment: Segment) -> Result<()> {
@@ -30,11 +30,19 @@ pub fn write_init(container: ContainerFormat, out: &mut Vec<u8>, segment: Segmen
         _ => unimplemented!(),
     }
 }
-pub fn write_frag(container: ContainerFormat, out: &mut Vec<u8>, segment: Segment) -> Result<()> {
+pub fn write_frag(
+    container: ContainerFormat,
+    out: &mut Vec<u8>,
+    segment: Segment,
+    index: usize,
+    end_pts: u64,
+) -> Result<()> {
     match container {
-        ContainerFormat::Matroska => MatroskaFragmentMuxer::write_frag(out, segment),
-        ContainerFormat::Webm => WebmFragmentMuxer::write_frag(out, segment),
-        ContainerFormat::MP4 => MP4FragmentMuxer::write_frag(out, segment),
+        ContainerFormat::Matroska => {
+            MatroskaFragmentMuxer::write_frag(out, segment, index, end_pts)
+        }
+        ContainerFormat::Webm => WebmFragmentMuxer::write_frag(out, segment, index, end_pts),
+        ContainerFormat::MP4 => MP4FragmentMuxer::write_frag(out, segment, index, end_pts),
         _ => unimplemented!(),
     }
 }
@@ -42,16 +50,19 @@ pub fn write_init_frag(
     container: ContainerFormat,
     out: &mut Vec<u8>,
     segment: Segment,
+    index: usize,
+    end_pts: u64,
 ) -> Result<()> {
     write_init(
         container,
         out,
         Segment {
             info: segment.info.clone(),
+            tags: segment.tags.clone(),
             tracks: segment.tracks.clone(),
             ..Default::default()
         },
     )?;
-    write_frag(container, out, segment)?;
+    write_frag(container, out, segment, index, end_pts)?;
     Ok(())
 }
diff --git a/remuxer/src/muxers/mp4.rs b/remuxer/src/muxers/mp4.rs
index bb89563..85bb244 100644
--- a/remuxer/src/muxers/mp4.rs
+++ b/remuxer/src/muxers/mp4.rs
@@ -4,6 +4,7 @@
     Copyright (C) 2026 metamuffin <metamuffin.org>
 */
 
+use std::sync::Mutex;
 use crate::muxers::FragmentMuxer;
 use anyhow::Result;
 use winter_matroska::{Segment, TrackType};
@@ -108,8 +109,8 @@ impl FragmentMuxer for MP4FragmentMuxer {
                 moov.write(MovieExtends(|mvex| {
                     mvex.write(MovieExtendsHeader {
                         fragment_duration: (segment.info.duration.unwrap_or_default()
-                            * segment.info.timestamp_scale as f64
-                            * 1e-9) as u64,
+                            * segment.info.timestamp_scale as f64)
+                            as u64,
                     });
                     for track in &segment.tracks.as_ref().unwrap().entries {
                         mvex.write(TrackExtends {
@@ -124,24 +125,67 @@ impl FragmentMuxer for MP4FragmentMuxer {
 
         Ok(())
     }
-    fn write_frag(out: &mut Vec<u8>, segment: Segment) -> Result<()> {
+    fn write_frag(out: &mut Vec<u8>, segment: Segment, index: usize, end_pts: u64) -> Result<()> {
         let mut out = BoxW::new(out);
 
+        let start_pts = (segment.clusters[0].timestamp as i64
+            + segment.clusters[0].simple_blocks[0].timestamp_off as i64)
+            * segment.info.timestamp_scale as i64;
+
+        let num_samples = segment
+            .clusters
+            .iter()
+            .map(|c| c.simple_blocks.len())
+            .sum::<usize>();
+        let dts_delta = ((end_pts - start_pts as u64) / num_samples as u64) as u32;
+        eprintln!("dts_delta {dts_delta}");
+
+        let mut mdat = Vec::new();
+        let mut samples = Vec::new();
+        for cluster in segment.clusters {
+            for block in cluster.simple_blocks {
+                let ts = cluster.timestamp as i64 + block.timestamp_off as i64;
+
+                samples.push(TrackRunSample {
+                    composition_time_offset: 0,
+                    duration: dts_delta as u32,
+                    flags: 0,
+                    size: block.data.len() as u32,
+                });
+
+                // samples.push((ts, block.data.len()));
+                mdat.extend(block.data);
+            }
+        }
+
         out.write(SegmentType {
             major_brand: *b"msdh",
             minor_version: 0,
             compatible_brands: &[*b"msdh", *b"msix"],
         });
+        let data_offset_offset = Mutex::new(0);
+        let base_data_offset = out.buffer_mut().len();
         out.write(MovieFragment(|moof| {
-            moof.write(MovieFragmentHeader { sequence_number: 0 });
+            moof.write(MovieFragmentHeader {
+                sequence_number: index as u32,
+            });
             moof.write(TrackFragment(|traf| {
-                traf.write(TrackFragmentHeader { track_id: 0 });
+                traf.write(TrackFragmentHeader { track_id: 1 });
                 traf.write(TrackFragmentBaseMediaDecodeTime {
                     base_media_decode_time: 0,
                 });
-                traf.write(TrackRun { data_offset: 0 });
+                *data_offset_offset.lock().unwrap() = traf.buffer_mut().len() + 12; // full box header (size+tag+ver+flags)
+                traf.write(TrackRun {
+                    data_offset: 0,
+                    samples: &samples,
+                });
             }));
         }));
+        let data_offset_offset = data_offset_offset.into_inner().unwrap();
+        let data_offset = out.buffer_mut().len() + 8 - base_data_offset; // box header (size+tag)
+        out.buffer_mut()[data_offset_offset..data_offset_offset + 4]
+            .copy_from_slice(&(data_offset as u32).to_be_bytes());
+        out.write(MediaData(&mdat));
 
         Ok(())
     }
diff --git a/stream/src/fragment.rs b/stream/src/fragment.rs
index ff50cf3..0ccf37d 100644
--- a/stream/src/fragment.rs
+++ b/stream/src/fragment.rs
@@ -62,6 +62,7 @@ pub fn fragment_init(
             track_num,
             file_track_num,
             0,
+            0,
             &iinfo.paths[file_index],
             format,
         )?;
@@ -146,6 +147,7 @@ pub fn fragment_transcode_full(
     track_num: TrackNum,
     file_track_num: u64,
     index: IndexNum,
+    end_pts: u64,
     media_path: &Path,
     output_format: &StreamFormatInfo,
 ) -> Result<Segment> {
@@ -155,6 +157,7 @@ pub fn fragment_transcode_full(
         .enumerate()
         .find(|(_, t)| t.remux)
         .unwrap();
+
     transcode(
         &sinfo.cache,
         &sinfo.config.transcoder,
@@ -162,12 +165,16 @@ pub fn fragment_transcode_full(
         &format!("{}-T{file_track_num}-I{index}", HashKey(&media_path)),
         input_format,
         output_format,
+        index,
+        end_pts,
         || {
             let init = fragment_init(&sinfo, track_num, input_format_num)?;
             let seg = fragment_remux(&sinfo, &media_path, file_track_num, index, true)?;
             Ok(Segment {
+                tracks: init.tracks,
+                info: init.info,
                 clusters: seg.clusters,
-                ..init
+                ..Default::default()
             })
         },
     )
@@ -196,6 +203,9 @@ pub fn fragment_stream(
         .get(format_num)
         .ok_or(anyhow!("format not found"))?;
 
+    let cue_stat = generate_cues(&sinfo.cache, &media_path)?;
+    let end_pts = cue_stat.cues.get(index + 1).map(|e| e.time).unwrap_or(0); // TODO
+
     let segment = if output_format.remux {
         fragment_remux(&sinfo, &media_path, file_track_num, index, false)?
     } else {
@@ -205,13 +215,14 @@ pub fn fragment_stream(
             track_num,
             file_track_num,
             index,
+            end_pts,
             &media_path,
             output_format,
         )?
     };
 
     let mut out = Vec::new();
-    write_frag(map_container(container), &mut out, segment)?;
+    write_frag(map_container(container), &mut out, segment, index, end_pts)?;
     Ok(Box::new(Cursor::new(out)))
 }
 
diff --git a/transcoder/src/fragment.rs b/transcoder/src/fragment.rs
index 58c8459..c34e8ee 100644
--- a/transcoder/src/fragment.rs
+++ b/transcoder/src/fragment.rs
@@ -62,6 +62,8 @@ pub fn transcode(
     input_key: &str,
     input_format: &StreamFormatInfo,
     output_format: &StreamFormatInfo,
+    index: usize,
+    end_pts: u64,
     segment: impl FnOnce() -> Result<Segment>,
 ) -> Result<Segment> {
     let command = transcode_command(kind, &input_format, output_format, false, config).unwrap();
@@ -86,7 +88,7 @@ pub fn transcode(
             let mut stdout = proc.stdout.take().unwrap();
 
             let mut buf = Vec::new();
-            write_init_frag(ContainerFormat::Matroska, &mut buf, input)?;
+            write_init_frag(ContainerFormat::Matroska, &mut buf, input, index, end_pts)?;
 
             spawn(move || {
                 stdin.write_all(&buf).unwrap();
@@ -106,6 +108,7 @@ pub fn transcode(
 
     let info = demuxer.info()?;
     let tracks = demuxer.tracks()?;
+    let tags = demuxer.tags()?;
     let mut clusters = Vec::new();
     while let Some((_, cluster)) = demuxer.read_cluster()? {
         clusters.push(cluster);
@@ -122,6 +125,7 @@ pub fn transcode(
     Ok(Segment {
         info,
         tracks,
+        tags: tags.into_iter().collect(),
         clusters,
         ..Default::default()
     })