Merge branch 'rewrite-stream'

1 files changed, 56 insertions, 53 deletions

diff --git a/remuxer/src/fragment.rs b/remuxer/src/fragment.rs
index 9fa68f3..0da1ed5 100644
--- a/remuxer/src/fragment.rs
+++ b/remuxer/src/fragment.rs
@@ -5,11 +5,12 @@
 */
 
 use crate::{
-    ebml_header, ebml_segment_info, ebml_track_entry, seek_index::get_seek_index,
+    ebml_header, ebml_segment_info, ebml_track_entry,
+    metadata::{matroska_metadata, MatroskaMetadata},
+    seek_index::get_seek_index,
     segment_extractor::SegmentExtractIter,
 };
 use anyhow::{anyhow, Context, Result};
-use jellybase::common::{LocalTrack, Node, SourceTrackKind};
 use jellymatroska::{read::EbmlReader, write::EbmlWriter, Master, MatroskaTag};
 use log::{debug, info};
 use std::{
@@ -19,34 +20,35 @@ use std::{
     path::Path,
 };
 
-const FRAGMENT_LENGTH: f64 = 2.;
+const FRAGMENT_LENGTH: f64 = 5.;
 
-pub fn fragment_index(
-    path_base: &Path,
-    item: &Node,
-    local_track: &LocalTrack,
-    track_index: usize,
-) -> Result<Vec<Range<f64>>> {
-    let media_info = item.media.as_ref().unwrap();
-    let source_path = path_base.join(&local_track.path);
-    let index = get_seek_index(&source_path)?;
+pub fn fragment_index(path: &Path, track: u64) -> Result<Vec<Range<f64>>> {
+    let meta = matroska_metadata(path)?;
+    let duration = media_duration(&meta);
+    let force_kf = meta
+        .as_ref()
+        .tracks
+        .as_ref()
+        .unwrap()
+        .entries
+        .iter()
+        .find(|t| t.track_number == track)
+        .unwrap()
+        .track_type
+        == 17;
+
+    let index = get_seek_index(&path)?;
     let index = index
-        .get(&(local_track.track as u64))
+        .get(&track)
         .ok_or(anyhow!("seek index track missing"))?;
 
-    // everything is a keyframe (even though nothing is...)
-    let force_kf = matches!(
-        media_info.tracks[track_index].kind,
-        SourceTrackKind::Subtitles { .. }
-    );
-
     let n_kf = if force_kf {
         index.blocks.len()
     } else {
         index.keyframes.len()
     };
 
-    let average_kf_interval = media_info.duration / n_kf as f64;
+    let average_kf_interval = duration / n_kf as f64;
     let kf_per_frag = (FRAGMENT_LENGTH / average_kf_interval).ceil() as usize;
     debug!("average keyframe interval: {average_kf_interval}");
     debug!(" => keyframes per frag {kf_per_frag}");
@@ -72,7 +74,7 @@ pub fn fragment_index(
                 index.keyframes.get((i + 1) * kf_per_frag).copied()
             }
             .map(|i| index.blocks[i].pts as f64 / 1000.)
-            .unwrap_or(media_info.duration);
+            .unwrap_or(duration);
             start..end
         })
         .collect())
@@ -80,45 +82,46 @@ pub fn fragment_index(
 
 pub fn write_fragment_into(
     writer: impl Write,
-    path_base: &Path,
-    item: &Node,
-    local_track: &LocalTrack,
-    track: usize,
+    path: &Path,
+    track: u64,
     webm: bool,
+    title: &str,
     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
+    let meta = matroska_metadata(path)?;
+    let duration = media_duration(&meta);
+    let track_meta = meta
+        .as_ref()
         .tracks
-        .get(track)
-        .ok_or(anyhow!("track not available"))?
-        .to_owned();
-    let source_path = path_base.join(&local_track.path);
+        .as_ref()
+        .unwrap()
+        .entries
+        .iter()
+        .find(|t| t.track_number == track)
+        .unwrap();
+    let force_kf = track_meta.track_type == 17;
+
+    info!("writing fragment {n} of {:?} (track {track})", title);
+    let mut output = EbmlWriter::new(BufWriter::new(writer), 0);
     let mapped = 1;
-    info!(
-        "\t- {track} {source_path:?} ({} => {mapped})",
-        local_track.track
-    );
-    info!("\t    {}", info);
-    let file = File::open(&source_path).context("opening source file")?;
-    let index = get_seek_index(&source_path)?;
+    info!("\t- {track} {path:?} ({} => {mapped})", track);
+    // info!("\t    {}", info);
+    let file = File::open(&path).context("opening source file")?;
+    let index = get_seek_index(&path)?;
     let index = index
-        .get(&(local_track.track as u64))
+        .get(&track)
         .ok_or(anyhow!("track missing 2"))?
         .to_owned();
     debug!("\t    seek index: {} blocks loaded", index.blocks.len());
     let mut reader = EbmlReader::new(BufReader::new(file));
 
-    let force_kf = matches!(info.kind, SourceTrackKind::Subtitles { .. });
     let n_kf = if force_kf {
         index.blocks.len()
     } else {
         index.keyframes.len()
     };
-
-    let average_kf_interval = media_info.duration / n_kf as f64;
+    debug!("{duration} {n_kf}");
+    let average_kf_interval = duration / n_kf as f64;
     let kf_per_frag = (FRAGMENT_LENGTH / average_kf_interval).ceil() as usize;
     debug!("average keyframe interval: {average_kf_interval}");
     debug!(" => keyframes per frag {kf_per_frag}");
@@ -144,25 +147,20 @@ pub fn write_fragment_into(
         .blocks
         .get(end_block_index)
         .map(|b| b.pts)
-        .unwrap_or((media_info.duration * 1000.) as u64);
+        .unwrap_or((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()),
+        title.to_string(),
         (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(),
-        ),
+        ebml_track_entry(mapped, track_meta),
     ])))?;
 
     reader.seek(start_block.source_off, MatroskaTag::Cluster(Master::Start))?;
-    let mut reader = SegmentExtractIter::new(&mut reader, local_track.track as u64);
+    let mut reader = SegmentExtractIter::new(&mut reader, track);
 
     {
         // TODO this one caused fragments to get dropped by MSE for no reason
@@ -214,3 +212,8 @@ pub fn write_fragment_into(
     debug!("wrote {} bytes", output.position());
     Ok(())
 }
+
+fn media_duration(m: &MatroskaMetadata) -> f64 {
+    let info = m.info.as_ref().unwrap();
+    (info.duration.unwrap_or_default() * info.timestamp_scale as f64) / 1_000_000_000.
+}