1 files changed, 150 insertions, 0 deletions
diff --git a/remuxer/src/seek_index.rs b/remuxer/src/seek_index.rs
new file mode 100644
index 0000000..7dbb9f7
--- /dev/null
+++ b/remuxer/src/seek_index.rs
@@ -0,0 +1,150 @@
+/*
+    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) 2023 metamuffin <metamuffin.org>
+*/
+use anyhow::Result;
+use jellycommon::{BlockIndex, SeekIndex};
+use jellymatroska::{
+    block::Block,
+    read::EbmlReader,
+    unflatten::{IterWithPos, Unflat, Unflatten},
+    MatroskaTag,
+};
+use log::{debug, info, trace, warn};
+use std::{collections::BTreeMap, fs::File, path::Path};
+
+pub fn write_all(path: &Path) -> Result<()> {
+    if path.with_extension(&format!("si.1")).exists() {
+        info!("seek index already present");
+        return Ok(());
+    }
+    let seek_index = {
+        let input = File::open(&path).unwrap();
+        let mut input = EbmlReader::new(input);
+        import_seek_index(&mut input)?
+    };
+    for (tn, index) in seek_index {
+        info!("writing index {tn} with {} blocks", index.blocks.len());
+        bincode::encode_into_std_write(
+            index,
+            &mut File::create(path.with_extension(&format!("si.{tn}")))?,
+            bincode::config::standard(),
+        )?;
+    }
+    Ok(())
+}
+
+pub fn import_seek_index(input: &mut EbmlReader) -> Result<BTreeMap<u64, SeekIndex>> {
+    let mut seek_index = BTreeMap::new();
+    while let Some(item) = input.next() {
+        let item = match item {
+            Ok(item) => item,
+            Err(e) => {
+                if !matches!(e, jellymatroska::error::Error::Io(_)) {
+                    warn!("{e}");
+                }
+                break;
+            }
+        };
+        match item {
+            MatroskaTag::Segment(_) => {
+                info!("segment start");
+                let mut children = Unflatten::new_with_end(input, item);
+                import_seek_index_segment(&mut children, &mut seek_index)?;
+                info!("segment end");
+            }
+            _ => debug!("(r) tag ignored: {item:?}"),
+        }
+    }
+    Ok(seek_index)
+}
+
+fn import_seek_index_segment(
+    segment: &mut Unflatten,
+    seek_index: &mut BTreeMap<u64, SeekIndex>,
+) -> Result<()> {
+    while let Some(Ok(Unflat { children, item, .. })) = segment.n() {
+        match item {
+            MatroskaTag::SeekHead(_) => {}
+            MatroskaTag::Info(_) => {}
+            MatroskaTag::Tags(_) => {}
+            MatroskaTag::Cues(_) => {}
+            MatroskaTag::Chapters(_) => {}
+            MatroskaTag::Tracks(_) => {}
+            MatroskaTag::Void(_) => {}
+            MatroskaTag::Cluster(_) => {
+                let mut children = children.unwrap();
+                let mut pts = 0;
+                let mut position = children.position();
+
+                loop {
+                    if let Some(Ok(Unflat { children, item, .. })) = children.n() {
+                        match item {
+                            MatroskaTag::Timestamp(ts) => pts = ts,
+                            MatroskaTag::BlockGroup(_) => {
+                                trace!("group");
+                                let mut children = children.unwrap();
+                                // let position = children.position(); //? TODO where should this point to? cluster or block? // probably block
+                                while let Some(Ok(Unflat {
+                                    children: _,
+                                    item,
+                                    position,
+                                })) = children.n()
+                                {
+                                    match item {
+                                        MatroskaTag::Block(ref buf) => {
+                                            let block = Block::parse(buf)?;
+                                            debug!(
+                                                "block: track={} tso={}",
+                                                block.track, block.timestamp_off
+                                            );
+                                            seek_index_add(seek_index, &block, position, pts);
+                                        }
+                                        _ => trace!("{item:?}"),
+                                    }
+                                }
+                            }
+                            MatroskaTag::SimpleBlock(buf) => {
+                                let block = Block::parse(&buf)?;
+                                trace!(
+                                    "simple block: track={} tso={}",
+                                    block.track,
+                                    block.timestamp_off
+                                );
+                                trace!("{pts} {}", block.timestamp_off);
+                                seek_index_add(seek_index, &block, position, pts);
+                            }
+                            _ => trace!("(rsc) tag ignored: {item:?}"),
+                        }
+                    } else {
+                        break;
+                    }
+                    position = children.position();
+                }
+            }
+            _ => debug!("(rs) tag ignored: {item:?}"),
+        };
+    }
+    Ok(())
+}
+
+fn seek_index_add(
+    seek_index: &mut BTreeMap<u64, SeekIndex>,
+    block: &Block,
+    position: usize,
+    pts_base: u64,
+) {
+    let trs = seek_index
+        .entry(block.track)
+        .or_insert(SeekIndex::default());
+
+    if block.keyframe {
+        trs.keyframes.push(trs.blocks.len());
+    }
+    trs.blocks.push(BlockIndex {
+        pts: pts_base + block.timestamp_off as u64,
+        source_off: position,
+        size: block.data.len(),
+    });
+}