consistent name for {snippet,segment?,fragment}

1 files changed, 212 insertions, 0 deletions

diff --git a/remuxer/src/fragment.rs b/remuxer/src/fragment.rs
new file mode 100644
index 0000000..660e2d9
--- /dev/null
+++ b/remuxer/src/fragment.rs
@@ -0,0 +1,212 @@
+/*
+    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) 2024 metamuffin <metamuffin.org>
+*/
+
+use crate::{
+    ebml_header, ebml_segment_info, ebml_track_entry, seek_index::get_seek_index,
+    segment_extractor::SegmentExtractIter,
+};
+use anyhow::{anyhow, Context, Result};
+use jellycommon::{LocalTrack, NodePublic, SourceTrackKind};
+use jellymatroska::{read::EbmlReader, write::EbmlWriter, Master, MatroskaTag};
+use log::{debug, info};
+use std::{
+    fs::File,
+    io::{BufReader, BufWriter, Write},
+    ops::Range,
+    path::Path,
+};
+
+const FRAGMENT_LENGTH: f64 = 2.;
+
+pub fn fragment_index(
+    path_base: &Path,
+    item: &NodePublic,
+    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)?;
+    let index = index
+        .get(&(local_track.track as u64))
+        .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 kf_per_snip = (FRAGMENT_LENGTH / average_kf_interval).ceil() as usize;
+    debug!("average keyframe interval: {average_kf_interval}");
+    debug!(" => keyframes per frag {kf_per_snip}");
+
+    let n_snips = n_kf.div_ceil(kf_per_snip);
+    Ok((0..n_snips)
+        .map(|i| {
+            let start = index.blocks[if force_kf {
+                i * kf_per_snip
+            } else {
+                index.keyframes[i * kf_per_snip]
+            }]
+            .pts as f64
+                / 1000.;
+            let end = if force_kf {
+                let n = (i + 1) * kf_per_snip;
+                if n >= index.blocks.len() {
+                    None
+                } else {
+                    Some(n)
+                }
+            } else {
+                index.keyframes.get((i + 1) * kf_per_snip).copied()
+            }
+            .map(|i| index.blocks[i].pts as f64 / 1000.)
+            .unwrap_or(media_info.duration);
+            start..end
+        })
+        .collect())
+}
+
+pub fn write_fragment_into(
+    writer: impl Write,
+    path_base: &Path,
+    item: &NodePublic,
+    local_track: &LocalTrack,
+    track: usize,
+    webm: bool,
+    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
+        .get(track)
+        .ok_or(anyhow!("track not available"))?
+        .to_owned();
+    let source_path = path_base.join(&local_track.path);
+    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)?;
+    let index = index
+        .get(&(local_track.track as u64))
+        .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;
+    let kf_per_snip = (FRAGMENT_LENGTH / average_kf_interval).ceil() as usize;
+    debug!("average keyframe interval: {average_kf_interval}");
+    debug!(" => keyframes per frag {kf_per_snip}");
+
+    let (start_block_index, end_block_index) = if force_kf {
+        (n * kf_per_snip, (n + 1) * kf_per_snip)
+    } else {
+        (
+            *index
+                .keyframes
+                .get(n * kf_per_snip)
+                .ok_or(anyhow!("fragment index out of range"))?,
+            *index
+                .keyframes
+                .get((n + 1) * kf_per_snip)
+                .unwrap_or(&index.blocks.len()),
+        )
+    };
+    debug!("writing blocks {start_block_index} to {end_block_index}.");
+
+    let start_block = &index.blocks[start_block_index];
+    let last_block_pts = index
+        .blocks
+        .get(end_block_index)
+        .map(|b| b.pts)
+        .unwrap_or((media_info.duration * 1000.) as u64);
+
+    reader.seek(start_block.source_off, MatroskaTag::Cluster(Master::Start))?;
+
+    output.write_tag(&ebml_header(webm))?;
+    output.write_tag(&MatroskaTag::Segment(Master::Start))?;
+    output.write_tag(&ebml_segment_info(
+        format!(
+            "{} (track {track}; frag {n})",
+            item.title.clone().unwrap_or_default()
+        ),
+        (last_block_pts - start_block.pts) as f64 / 1000.,
+    ))?;
+    output.write_tag(&MatroskaTag::Tags(Master::Collected(vec![])))?;
+    output.write_tag(&MatroskaTag::Tracks(Master::Collected(vec![
+        ebml_track_entry(mapped, &info, local_track.codec_private.clone()),
+    ])))?;
+
+    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 = reader.next()?;
+
+        //     assert_eq!(index_block.size, block.data.len(), "seek index is wrong");
+
+        //     block.track = 1;
+        //     block.timestamp_off = 0;
+        //     output.write_tag(&MatroskaTag::Cluster(Master::Collected(vec![
+        //         MatroskaTag::Timestamp(index_block.pts - start_block.pts),
+        //         MatroskaTag::SimpleBlock(block.dump()),
+        //     ])))?;
+        // }
+    }
+    {
+        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()?;
+
+            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 {
+                blocks.push(MatroskaTag::BlockGroup(Master::Collected(vec![
+                    MatroskaTag::BlockDuration(duration),
+                    MatroskaTag::Block(block),
+                ])))
+            } else {
+                blocks.push(MatroskaTag::SimpleBlock(block))
+            }
+        }
+        output.write_tag(&MatroskaTag::Cluster(Master::Collected(blocks)))?;
+    }
+
+    // output.write_tag(&MatroskaTag::Segment(Master::End))?;
+
+    debug!("wrote {} bytes", output.position());
+    Ok(())
+}