non-seekable mkv almost works

5 files changed, 228 insertions, 84 deletions

diff --git a/common/src/lib.rs b/common/src/lib.rs
index c4eac09..0eacc42 100644
--- a/common/src/lib.rs
+++ b/common/src/lib.rs
@@ -12,7 +12,7 @@ pub struct DirectoryInfo {
 pub struct ItemInfo {
     pub title: String,
     pub duration: f64, // in seconds
-    pub tracks: BTreeMap<usize, SourceTrack>,
+    pub tracks: BTreeMap<u64, SourceTrack>,
 }
 
 #[derive(Debug, Clone, Deserialize, Serialize)]
@@ -21,6 +21,7 @@ pub struct SourceTrack {
     pub name: String,
     pub codec: String,
     pub language: String,
+    pub codec_private: Option<Vec<u8>>,
 }
 
 #[derive(Debug, Clone, Deserialize, Serialize)]
diff --git a/matroska/src/block.rs b/matroska/src/block.rs
index b0d6f6b..dd5b340 100644
--- a/matroska/src/block.rs
+++ b/matroska/src/block.rs
@@ -1,4 +1,4 @@
-use crate::read::ReadExt;
+use crate::{read::ReadExt, write::write_vint};
 use anyhow::Result;
 use std::io::Cursor;
 
@@ -40,4 +40,21 @@ impl Block {
             timestamp_off,
         })
     }
+    pub fn dump(&self) -> Vec<u8> {
+        let mut out = vec![];
+        write_vint(&mut out, self.track).unwrap();
+        out.extend(self.timestamp_off.to_be_bytes().into_iter());
+        out.push(
+            match self.invisible {
+                true => 0b10000,
+                false => 0b00000,
+            } | match self.lacing {
+                Some(LacingType::Xiph) => 0b0100,
+                Some(LacingType::Ebml) => 0b1000,
+                Some(LacingType::FixedSize) => 0b1100,
+                None => 0b0000,
+            },
+        );
+        out
+    }
 }
diff --git a/remuxer/src/import/mod.rs b/remuxer/src/import/mod.rs
index bd22cf5..11970b0 100644
--- a/remuxer/src/import/mod.rs
+++ b/remuxer/src/import/mod.rs
@@ -1,8 +1,7 @@
-use std::collections::HashMap;
-
 use anyhow::{anyhow, bail, Result};
 use jellycommon::{ItemInfo, SourceTrack, SourceTrackKind};
 use jellymatroska::{
+    block::Block,
     matroska::MatroskaTag,
     read::EbmlReader,
     unflatten::{Unflat, Unflatten},
@@ -47,7 +46,7 @@ pub fn import_read(input: &mut EbmlReader, iteminfo: &mut ItemInfo) -> Result<()
 }
 
 fn import_read_segment(children: &mut Unflatten, iteminfo: &mut ItemInfo) -> Result<()> {
-    let mut track_mapping = HashMap::<u64, usize>::new(); // maps matroska track id to item track id
+    // let mut track_mapping = HashMap::<u64, usize>::new(); // maps matroska track id to item track id
     let (mut timestamp_scale, mut duration) = (None, None);
     while let Some(Ok(Unflat { children, item })) = children.next() {
         match item {
@@ -72,13 +71,23 @@ fn import_read_segment(children: &mut Unflatten, iteminfo: &mut ItemInfo) -> Res
                             let mut children = children.unwrap();
                             while let Some(Ok(Unflat { children, item })) = children.next() {
                                 match item {
-                                    MatroskaTag::Block(_) => (),
+                                    MatroskaTag::Block(buf) => {
+                                        let block = Block::parse(&buf)?;
+                                        debug!(
+                                            "block: track={} tso={}",
+                                            block.track, block.timestamp_off
+                                        )
+                                    }
                                     _ => trace!("{item:?}"),
                                 }
                             }
                         }
-                        MatroskaTag::SimpleBlock(_) => {
-                            // debug!("simple");
+                        MatroskaTag::SimpleBlock(buf) => {
+                            let block = Block::parse(&buf)?;
+                            debug!(
+                                "simple block: track={} tso={}",
+                                block.track, block.timestamp_off
+                            )
                         }
                         _ => debug!("(rsc) tag ignored: {item:?}"),
                     }
@@ -118,8 +127,10 @@ fn import_read_segment(children: &mut Unflatten, iteminfo: &mut ItemInfo) -> Res
                                 mut name,
                                 mut fps,
                                 mut bit_depth,
+                                mut codec_private,
                             ) = (
                                 None, None, None, None, None, None, None, None, None, None, None,
+                                None,
                             );
                             while let Some(Ok(Unflat { children, item })) = children.next() {
                                 match item {
@@ -128,6 +139,7 @@ fn import_read_segment(children: &mut Unflatten, iteminfo: &mut ItemInfo) -> Res
                                     MatroskaTag::TrackNumber(v) => index = Some(v),
                                     MatroskaTag::TrackType(v) => kind = Some(v),
                                     MatroskaTag::Name(v) => name = Some(v),
+                                    MatroskaTag::CodecPrivate(v) => codec_private = Some(v),
                                     MatroskaTag::Audio(_) => {
                                         let mut children = children.unwrap();
                                         while let Some(Ok(Unflat { item, .. })) = children.next() {
@@ -157,11 +169,11 @@ fn import_read_segment(children: &mut Unflatten, iteminfo: &mut ItemInfo) -> Res
                                     _ => (),
                                 }
                             }
-                            let itrack_index = iteminfo.tracks.len();
+                            // let itrack_index = iteminfo.tracks.len();
                             let mtrack_index = index.unwrap();
                             let kind = match kind.ok_or(anyhow!("track type required"))? {
                                 1 => SourceTrackKind::Video {
-                                    fps: fps.unwrap_or(f64::NAN), // TODO
+                                    fps: fps.unwrap_or(0.0), // TODO
                                     width: width.unwrap(),
                                     height: height.unwrap(),
                                 },
@@ -173,10 +185,11 @@ fn import_read_segment(children: &mut Unflatten, iteminfo: &mut ItemInfo) -> Res
                                 17 => SourceTrackKind::Subtitles,
                                 _ => bail!("invalid track type"),
                             };
-                            track_mapping.insert(mtrack_index, itrack_index);
+                            // track_mapping.insert(mtrack_index, itrack_index);
                             iteminfo.tracks.insert(
-                                itrack_index,
+                                mtrack_index,
                                 SourceTrack {
+                                    codec_private,
                                     name: name.unwrap_or_else(|| "unnamed".to_string()),
                                     codec: codec.unwrap(),
                                     language: language.unwrap_or_else(|| "none".to_string()),
diff --git a/remuxer/src/lib.rs b/remuxer/src/lib.rs
index cc508e8..523569a 100644
--- a/remuxer/src/lib.rs
+++ b/remuxer/src/lib.rs
@@ -1,8 +1,23 @@
-pub mod import;
 pub mod format;
+pub mod import;
 
-use jellycommon::ItemInfo;
-use std::{io::Write, path::PathBuf, sync::Arc};
+use anyhow::Result;
+use jellycommon::{ItemInfo, SourceTrack, SourceTrackKind};
+use jellymatroska::{
+    block::Block,
+    read::EbmlReader,
+    unflatten::{Unflat, Unflatten},
+    write::{write_vint, EbmlWriter},
+    Master, MatroskaTag,
+};
+use log::{debug, error, info, trace, warn};
+use std::{
+    collections::{BTreeMap, HashMap},
+    fs::File,
+    io::Write,
+    path::PathBuf,
+    sync::Arc,
+};
 
 pub struct RemuxerContext {}
 
@@ -13,83 +28,179 @@ impl RemuxerContext {
 
     pub fn generate_into(
         &self,
-        writer: impl Write,
+        writer: impl Write + 'static,
         offset: usize,
         path_base: PathBuf,
-        item: ItemInfo,
+        iteminfo: ItemInfo,
         selection: Vec<u64>,
     ) -> anyhow::Result<()> {
-        // let source_path = path_base.join(item.source.path);
-        // info!("remuxing {source_path:?} to have tracks {selection:?}");
-        // let mut input = File::open(source_path)?;
+        let source_path = path_base.join(format!("demon-slayer-1.mkv"));
+        info!("remuxing {source_path:?} to have tracks {selection:?}");
+
+        let input = File::open(source_path)?;
+        let mut input = EbmlReader::new(input);
+        let mut output = EbmlWriter::new(writer, 0);
+
+        // maps original to remuxed
+        let mapping = BTreeMap::from_iter(
+            selection
+                .iter()
+                .enumerate()
+                .map(|(i, e)| (*e, i as u64 + 1)),
+        );
+        info!("track mapping: {mapping:?}");
+
+        output.write_tag(&MatroskaTag::Ebml(Master::Collected(vec![
+            MatroskaTag::EbmlVersion(1),
+            MatroskaTag::EbmlReadVersion(1),
+            MatroskaTag::EbmlMaxIdLength(4),
+            MatroskaTag::EbmlMaxSizeLength(8),
+            MatroskaTag::DocType("matroska".to_string()),
+            MatroskaTag::DocTypeVersion(4),
+            MatroskaTag::DocTypeReadVersion(2),
+        ])))?;
+
+        output.write_tag(&MatroskaTag::Segment(Master::Start))?;
+
+        output.write_tag(&MatroskaTag::Info(Master::Collected(vec![
+            MatroskaTag::Title(iteminfo.title.clone()),
+            MatroskaTag::Duration(1000.0),
+            MatroskaTag::MuxingApp("jellyremux".to_string()),
+            MatroskaTag::WritingApp("jellything".to_string()),
+        ])))?;
+
+        let tracks_header = MatroskaTag::Tracks(Master::Collected(
+            mapping
+                .iter()
+                .map(|(norig, nrem)| track_to_ebml(*nrem, &iteminfo.tracks.get(norig).unwrap()))
+                .collect(),
+        ));
+        info!("track header: {tracks_header:?}");
+        output.write_tag(&tracks_header)?;
 
-        // let tags = WebmIterator::new(&mut input, &[MatroskaSpec::TrackEntry(Master::Start)]);
-        // let mut output = WebmWriter::new(writer);
+        while let Some(item) = input.next() {
+            let item = match item {
+                Ok(item) => item,
+                Err(e) => {
+                    warn!("{e}");
+                    break;
+                }
+            };
+            match item {
+                MatroskaTag::Segment(_) => {
+                    info!("segment start");
+                    let mut children = Unflatten::new_with_end(&mut input, item);
+                    filter_segment(&mut children, &mut output, &iteminfo, &mapping)?;
+                    info!("segment end");
+                }
+                _ => debug!("(r) tag ignored: {item:?}"),
+            }
+        }
 
-        // let mut tscale = None;
-        // let mut duration = None;
-        // let mut ignore = false;
+        output.write_tag(&MatroskaTag::Segment(Master::End))?;
 
-        // for tag in tags {
-        //     let tag = tag.unwrap();
-        //     match tag {
-        //         MatroskaSpec::SeekHead(Master::Start) | MatroskaSpec::Cues(Master::Start) => {
-        //             ignore = true
-        //         }
-        //         MatroskaSpec::SeekHead(Master::End) | MatroskaSpec::Cues(Master::End) => {
-        //             ignore = false
-        //         }
-        //         MatroskaSpec::TrackEntry(master) => {
-        //             let children = master.get_children();
-        //             let mut number = None;
-        //             for c in &children {
-        //                 if let MatroskaSpec::TrackNumber(n) = c {
-        //                     number = Some(*n)
-        //                 }
-        //             }
-        //             let number = number.unwrap();
-        //             if selection.contains(&number) {
-        //                 output.write(&MatroskaSpec::TrackEntry(Master::Full(children)))?;
-        //             }
-        //         }
-        //         MatroskaSpec::Block(ref data) => {
-        //             let data: &[u8] = &data;
-        //             let block: Block = data.try_into()?;
-        //             if selection.contains(&block.track) {
-        //                 output.write(&tag)?;
-        //             }
-        //         }
-        //         MatroskaSpec::SimpleBlock(ref data) => {
-        //             let data: &[u8] = &data;
-        //             let block: Block = data.try_into()?;
-        //             if selection.contains(&block.track) {
-        //                 output.write(&tag)?;
-        //             }
-        //         }
-        //         MatroskaSpec::Info(Master::Start) => (),
-        //         MatroskaSpec::TimestampScale(n) => tscale = Some(n),
-        //         MatroskaSpec::Duration(n) => duration = Some(n),
-        //         MatroskaSpec::Info(Master::End) => {
-        //             output.write(&MatroskaSpec::Info(Master::Full(vec![
-        //                 MatroskaSpec::TimestampScale(tscale.unwrap()),
-        //                 MatroskaSpec::Title(item.title.clone()),
-        //                 MatroskaSpec::Duration(duration.unwrap()),
-        //                 MatroskaSpec::MuxingApp("jellyremux".to_string()),
-        //                 MatroskaSpec::WritingApp("jellything".to_string()),
-        //             ])))?;
-        //         }
-        //         x => {
-        //             if !ignore {
-        //                 debug!("{x:?}");
-        //                 output.write(&x)?;
-        //             }
-        //         }
-        //     }
-        // }
         Ok(())
     }
 }
 
+fn filter_segment(
+    children: &mut Unflatten,
+    writer: &mut EbmlWriter,
+    iteminfo: &ItemInfo,
+    mapping: &BTreeMap<u64, u64>,
+) -> Result<()> {
+    while let Some(Ok(Unflat { children, item })) = children.next() {
+        match item {
+            MatroskaTag::SeekHead(_) => {}
+            MatroskaTag::Info(_) => {}
+            MatroskaTag::Cluster(_) => {
+                let mut cluster = vec![];
+                info!("start of cluster found");
+                let mut children = children.unwrap();
+                while let Some(Ok(Unflat { children, item })) = children.next() {
+                    match item {
+                        MatroskaTag::BlockGroup(_) => {
+                            debug!("group");
+                            let mut children = children.unwrap();
+                            let mut group = vec![];
+                            while let Some(Ok(Unflat { children, item })) = children.next() {
+                                match item {
+                                    MatroskaTag::Block(buf) => {
+                                        let mut block = Block::parse(&buf)?;
+                                        if let Some(outnum) = mapping.get(&block.track) {
+                                            block.track = *outnum;
+                                            debug!(
+                                                "block: track={} tso={}",
+                                                block.track, block.timestamp_off
+                                            );
+                                            group.push(MatroskaTag::Block(block.dump()));
+                                        }
+                                    }
+                                    _ => trace!("{item:?}"),
+                                }
+                            }
+                            cluster.push(MatroskaTag::BlockGroup(Master::Collected(group)));
+                        }
+                        MatroskaTag::SimpleBlock(buf) => {
+                            let mut block = Block::parse(&buf)?;
+                            if let Some(outnum) = mapping.get(&block.track) {
+                                block.track = *outnum;
+                                debug!("block: track={} tso={}", block.track, block.timestamp_off);
+                                cluster.push(MatroskaTag::SimpleBlock(block.dump()));
+                            }
+                        }
+                        _ => debug!("(rsc) tag ignored: {item:?}"),
+                    }
+                }
+                writer.write_tag(&MatroskaTag::Cluster(Master::Collected(cluster)))?;
+            }
+            MatroskaTag::Tags(_) => {}
+            MatroskaTag::Cues(_) => {}
+            MatroskaTag::Chapters(_) => {}
+            MatroskaTag::Tracks(_) => {}
+            _ => debug!("(rs) tag ignored: {item:?}"),
+        }
+    }
+    Ok(())
+}
+
+pub fn track_to_ebml(number: u64, track: &SourceTrack) -> MatroskaTag {
+    let mut els = Vec::new();
+    els.push(MatroskaTag::TrackNumber(number));
+    els.push(MatroskaTag::TrackUID(number));
+    els.push(MatroskaTag::FlagLacing(0));
+    els.push(MatroskaTag::Language(track.language.clone()));
+    els.push(MatroskaTag::CodecID(track.codec.clone()));
+    match track.kind {
+        SourceTrackKind::Video { width, height, fps } => {
+            els.push(MatroskaTag::TrackType(1));
+            els.push(MatroskaTag::DefaultDuration(41708333)); // TODO pls not hardcoded
+            els.push(MatroskaTag::Video(Master::Collected(vec![
+                MatroskaTag::PixelWidth(width),
+                MatroskaTag::PixelHeight(height),
+            ])))
+        }
+        SourceTrackKind::Audio {
+            channels,
+            sample_rate,
+            bit_depth,
+        } => {
+            els.push(MatroskaTag::TrackType(2));
+            els.push(MatroskaTag::Audio(Master::Collected(vec![
+                MatroskaTag::SamplingFrequency(sample_rate),
+                MatroskaTag::Channels(channels.try_into().unwrap()),
+            ])))
+        }
+        SourceTrackKind::Subtitles => {
+            els.push(MatroskaTag::TrackType(19));
+        }
+    }
+    if let Some(d) = &track.codec_private {
+        els.push(MatroskaTag::CodecPrivate(d.clone()));
+    }
+    MatroskaTag::TrackEntry(Master::Collected(els))
+}
+
 // pub struct SendWriter(pub Sender<Vec<u8>>);
 
 // impl Write for SendWriter {
diff --git a/server/src/library.rs b/server/src/library.rs
index a59635b..c75bf69 100644
--- a/server/src/library.rs
+++ b/server/src/library.rs
@@ -1,5 +1,6 @@
 use anyhow::{anyhow, bail, Context, Ok};
 use jellycommon::{DirectoryInfo, ItemInfo};
+use log::info;
 use std::{ffi::OsStr, fs::File, path::PathBuf, sync::Arc};
 
 pub struct Library {
@@ -93,9 +94,9 @@ impl Node {
                 .read_dir()?
                 .filter_map(|e| {
                     let e = e.unwrap();
-                    if (e.path().extension() == Some(OsStr::new("json"))
-                        && e.path().file_name() != Some(OsStr::new("directory.json")))
-                        || e.metadata().unwrap().is_dir()
+                    if (e.path().extension() != Some(OsStr::new("mkv"))
+                        || e.metadata().unwrap().is_dir())
+                        && !e.path().ends_with("directory.json")
                     {
                         Some(e.path())
                     } else {
@@ -117,6 +118,7 @@ impl Node {
             }))
             .into())
         } else if path.is_file() {
+            info!("loading {path:?}");
             let datafile = File::open(path.clone()).context("cant load metadata")?;
             let data: ItemInfo = serde_json::from_reader(datafile).context("invalid metadata")?;
             let identifier = path