ebml-struct based fragment writer (does not work)

4 files changed, 215 insertions, 59 deletions

diff --git a/remuxer/Cargo.toml b/remuxer/Cargo.toml
index 2313dcc..acbdb44 100644
--- a/remuxer/Cargo.toml
+++ b/remuxer/Cargo.toml
@@ -13,3 +13,8 @@ log = { workspace = true }
 
 serde = { version = "1.0.217", features = ["derive"] }
 bincode = { version = "2.0.0-rc.3", features = ["serde"] }
+
+# ebml-struct = { git = "https://codeberg.org/metamuffin/ebml-struct", features = [
+#     "bincode",
+# ] }
+ebml-struct = { path = "../../ebml-struct", features = ["bincode"] }
diff --git a/remuxer/src/fragment.rs b/remuxer/src/fragment.rs
index 9fa68f3..c7954b0 100644
--- a/remuxer/src/fragment.rs
+++ b/remuxer/src/fragment.rs
@@ -5,12 +5,16 @@
 */
 
 use crate::{
-    ebml_header, ebml_segment_info, ebml_track_entry, seek_index::get_seek_index,
-    segment_extractor::SegmentExtractIter,
+    metadata::matroska_metadata, seek_index::get_seek_index, segment_extractor::SegmentExtractIter,
 };
 use anyhow::{anyhow, Context, Result};
+use ebml_struct::{
+    matroska::{BlockGroup, Cluster, Ebml, Info, Segment, Tracks},
+    write::TagWrite,
+    Block,
+};
 use jellybase::common::{LocalTrack, Node, SourceTrackKind};
-use jellymatroska::{read::EbmlReader, write::EbmlWriter, Master, MatroskaTag};
+use jellymatroska::{read::EbmlReader, Master, MatroskaTag};
 use log::{debug, info};
 use std::{
     fs::File,
@@ -88,7 +92,7 @@ pub fn write_fragment_into(
     n: usize,
 ) -> anyhow::Result<()> {
     info!("writing fragment {n} of {:?} (track {track})", item.title);
-    let mut output = EbmlWriter::new(BufWriter::new(writer), 0);
+
     let media_info = item.media.as_ref().unwrap();
     let info = media_info
         .tracks
@@ -146,71 +150,101 @@ pub fn write_fragment_into(
         .map(|b| b.pts)
         .unwrap_or((media_info.duration * 1000.) as u64);
 
-    output.write_tag(&ebml_header(webm))?;
-    output.write_tag(&MatroskaTag::Segment(Master::Start))?;
-    output.write_tag(&ebml_segment_info(
-        format!("{}: {info}", item.title.clone().unwrap_or_default()),
-        (last_block_pts - start_block.pts) as f64 / 1000.,
-    ))?;
-    output.write_tag(&MatroskaTag::Tracks(Master::Collected(vec![
-        ebml_track_entry(
-            mapped,
-            local_track.track as u64 * 100, // TODO something else that is unique to the track
-            &info,
-            local_track.codec_private.clone(),
-        ),
-    ])))?;
+    let input_metadata = (*matroska_metadata(&local_track.path)?).clone().unwrap();
 
     reader.seek(start_block.source_off, MatroskaTag::Cluster(Master::Start))?;
     let mut reader = SegmentExtractIter::new(&mut reader, local_track.track as u64);
 
-    {
-        // TODO this one caused fragments to get dropped by MSE for no reason
-        // for i in start_block_index..end_block_index {
-        //     let index_block = &index.blocks[i];
-        //     let (mut block, duration) = reader.next()?;
+    let mut cluster = Cluster::default();
+    cluster.timestamp = start_block.pts;
+    for i in start_block_index..end_block_index {
+        let index_block = &index.blocks[i];
+        let (block, duration) = reader.next_block()?;
 
-        //     assert_eq!(index_block.size, block.data.len(), "seek index is wrong");
+        let mut block = Block {
+            data: block.data,
+            discardable: block.discardable,
+            invisible: block.invisible,
+            keyframe: block.keyframe,
+            lacing: block.lacing.map(|l| match l {
+                jellymatroska::block::LacingType::Xiph => ebml_struct::LacingType::Xiph,
+                jellymatroska::block::LacingType::FixedSize => ebml_struct::LacingType::FixedSize,
+                jellymatroska::block::LacingType::Ebml => ebml_struct::LacingType::Ebml,
+            }),
+            timestamp_off: block.timestamp_off,
+            track: block.track,
+        };
 
-        //     block.track = 1;
-        //     block.timestamp_off = 0;
-        //     output.write_tag(&MatroskaTag::Cluster(Master::Collected(vec![
-        //         MatroskaTag::Timestamp(index_block.pts - start_block.pts),
-        //         if let Some(duration) = duration {
-        //             MatroskaTag::BlockGroup(Master::Collected(vec![
-        //                 MatroskaTag::BlockDuration(duration),
-        //                 MatroskaTag::Block(block),
-        //             ]))
-        //         } else {
-        //             MatroskaTag::SimpleBlock(block)
-        //         },
-        //     ])))?;
-        // }
+        assert_eq!(index_block.size, block.data.len(), "seek index is wrong");
+
+        block.track = 1;
+        // TODO this does generate overflows sometimes
+        block.timestamp_off = (index_block.pts as i64 - start_block.pts as i64)
+            .try_into()
+            .unwrap();
+        if let Some(duration) = duration {
+            cluster.block_groups.push(BlockGroup {
+                block_duration: Some(duration),
+                block,
+                ..Default::default()
+            })
+        } else {
+            cluster.simple_blocks.push(block)
+        }
     }
-    {
-        let mut blocks = vec![MatroskaTag::Timestamp(start_block.pts)];
-        for i in start_block_index..end_block_index {
-            let index_block = &index.blocks[i];
-            let (mut block, duration) = reader.next_block()?;
 
-            assert_eq!(index_block.size, block.data.len(), "seek index is wrong");
+    let mut input_track = input_metadata
+        .tracks
+        .unwrap()
+        .entries
+        .into_iter()
+        .find(|t| t.track_number == local_track.track as u64)
+        .unwrap();
 
-            block.track = 1;
-            // TODO this does generate overflows sometimes
-            block.timestamp_off = (index_block.pts as i64 - start_block.pts as i64)
-                .try_into()
-                .unwrap();
-            if let Some(duration) = duration {
-                blocks.push(MatroskaTag::BlockGroup(Master::Collected(vec![
-                    MatroskaTag::BlockDuration(duration),
-                    MatroskaTag::Block(block),
-                ])))
-            } else {
-                blocks.push(MatroskaTag::SimpleBlock(block))
-            }
+    input_track.track_number = 1;
+    if webm {
+        if let Some(v) = &mut input_track.video {
+            v.colour = None;
         }
-        output.write_tag(&MatroskaTag::Cluster(Master::Collected(blocks)))?;
     }
-    debug!("wrote {} bytes", output.position());
+
+    let mut output = BufWriter::new(writer);
+
+    Ebml {
+        ebml_version: 1,
+        ebml_read_version: 1,
+        ebml_max_id_length: 4,
+        ebml_max_size_length: 8,
+        doc_type: if webm {
+            "webm".to_string()
+        } else {
+            "matroska".to_string()
+        },
+        doc_type_version: 4,
+        doc_type_read_version: 2,
+        doc_type_extensions: vec![],
+    }
+    .write(&mut output)?;
+
+    Segment {
+        info: Info {
+            timestamp_scale: 1_000_000,
+            duration: Some((last_block_pts - start_block.pts) as f64),
+            title: Some(format!(
+                "{}: {info}",
+                item.title.clone().unwrap_or_default()
+            )),
+            muxing_app: "ebml-struct".to_owned(),
+            writing_app: "jellything".to_owned(),
+            ..Default::default()
+        },
+        tracks: Some(Tracks {
+            entries: vec![input_track],
+        }),
+        clusters: vec![cluster],
+        ..Default::default()
+    }
+    .write(&mut output)?;
+
     Ok(())
 }
diff --git a/remuxer/src/lib.rs b/remuxer/src/lib.rs
index b46369e..f3526dc 100644
--- a/remuxer/src/lib.rs
+++ b/remuxer/src/lib.rs
@@ -9,6 +9,7 @@ pub mod remux;
 pub mod seek_index;
 pub mod segment_extractor;
 pub mod trim_writer;
+pub mod metadata;
 
 pub use fragment::write_fragment_into;
 pub use remux::remux_stream_into;
diff --git a/remuxer/src/metadata.rs b/remuxer/src/metadata.rs
new file mode 100644
index 0000000..92b5445
--- /dev/null
+++ b/remuxer/src/metadata.rs
@@ -0,0 +1,116 @@
+/*
+    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) 2025 metamuffin <metamuffin.org>
+*/
+use anyhow::{Context, Result};
+use bincode::{Decode, Encode};
+use ebml_struct::{
+    ids::*,
+    matroska::*,
+    read::{EbmlReadExt, TagRead},
+};
+use jellybase::{
+    assetfed::AssetInner,
+    cache::{cache_file, cache_memory},
+    common::Asset,
+};
+use log::info;
+use std::{
+    fs::File,
+    io::{BufReader, ErrorKind, Read, Write},
+    path::Path,
+    sync::Arc,
+};
+
+#[derive(Encode, Decode, Clone)]
+pub struct MatroskaMetadata {
+    pub info: Option<Info>,
+    pub tracks: Option<Tracks>,
+    pub cover: Option<Asset>,
+    pub chapters: Option<Chapters>,
+    pub tags: Option<Tags>,
+    pub infojson: Option<Vec<u8>>,
+}
+pub fn matroska_metadata(path: &Path) -> Result<Arc<Option<MatroskaMetadata>>> {
+    cache_memory(&["mkmeta-v1", path.to_string_lossy().as_ref()], || {
+        let mut magic = [0; 4];
+        File::open(path)?.read_exact(&mut magic).ok();
+        if !matches!(magic, [0x1A, 0x45, 0xDF, 0xA3]) {
+            return Ok(None);
+        }
+
+        info!("reading media file {path:?}");
+        let mut file = BufReader::new(File::open(path)?);
+        let mut file = file.by_ref().take(u64::MAX);
+
+        let (x, mut ebml) = file.read_tag()?;
+        assert_eq!(x, EL_EBML);
+        let ebml = Ebml::read(&mut ebml).unwrap();
+        assert!(ebml.doc_type == "matroska" || ebml.doc_type == "webm");
+        let (x, mut segment) = file.read_tag()?;
+        assert_eq!(x, EL_SEGMENT);
+
+        let mut info = None;
+        let mut infojson = None;
+        let mut tracks = None;
+        let mut cover = None;
+        let mut chapters = None;
+        let mut tags = None;
+        loop {
+            let (x, mut seg) = match segment.read_tag() {
+                Ok(o) => o,
+                Err(e) if e.kind() == ErrorKind::UnexpectedEof => break,
+                Err(e) => return Err(e.into()),
+            };
+            match x {
+                EL_INFO => info = Some(Info::read(&mut seg).context("info")?),
+                EL_TRACKS => tracks = Some(Tracks::read(&mut seg).context("tracks")?),
+                EL_CHAPTERS => chapters = Some(Chapters::read(&mut seg).context("chapters")?),
+                EL_TAGS => tags = Some(Tags::read(&mut seg).context("tags")?),
+                EL_ATTACHMENTS => {
+                    let attachments = Attachments::read(&mut seg).context("attachments")?;
+                    for f in attachments.files {
+                        match f.name.as_str() {
+                            "info.json" => {
+                                infojson = Some(f.data);
+                            }
+                            "cover.webp" | "cover.png" | "cover.jpg" | "cover.jpeg"
+                            | "cover.avif" => {
+                                cover = Some(
+                                    AssetInner::Cache(cache_file(
+                                        &["att-cover", path.to_string_lossy().as_ref()],
+                                        move |mut file| {
+                                            file.write_all(&f.data)?;
+                                            Ok(())
+                                        },
+                                    )?)
+                                    .ser(),
+                                )
+                            }
+                            a => println!("{a:?}"),
+                        }
+                    }
+                }
+                EL_VOID | EL_CRC32 | EL_CUES | EL_SEEKHEAD => {
+                    seg.consume()?;
+                }
+                EL_CLUSTER => {
+                    break;
+                }
+                id => {
+                    eprintln!("unknown top-level element {id:x}");
+                    seg.consume()?;
+                }
+            }
+        }
+        Ok(Some(MatroskaMetadata {
+            chapters,
+            cover,
+            info,
+            infojson,
+            tags,
+            tracks,
+        }))
+    })
+}