/* 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 */ use anyhow::{Context, Result}; use bincode::{Decode, Encode}; use jellybase::cache::cache_memory; use jellymatroska::{ block::Block, read::EbmlReader, unflatten::{Unflat, Unflatten}, MatroskaTag, }; use log::{debug, info, trace, warn}; use std::{collections::BTreeMap, fs::File, io::BufReader, path::Path, sync::Arc}; #[derive(Debug, Clone, Decode, Encode)] pub struct SeekIndex { pub blocks: Vec, pub keyframes: Vec, } #[derive(Debug, Clone, Decode, Encode)] pub struct BlockIndex { pub pts: u64, // pub duration: Option, pub source_off: u64, // points to start of SimpleBlock or BlockGroup (not the Block inside it) pub size: usize, } impl Default for SeekIndex { fn default() -> Self { Self { blocks: Vec::new(), keyframes: Vec::new(), } } } pub fn get_seek_index(path: &Path) -> anyhow::Result>>> { cache_memory("seekindex-v1", path, move || { info!("generating seek index for {path:?}"); let input = File::open(path).context("opening source file")?; let mut input = EbmlReader::new(BufReader::new(input)); let index = import_seek_index(&mut input)?; info!("done"); Ok(index.into_iter().map(|(k, v)| (k, Arc::new(v))).collect()) }) } pub fn get_track_sizes(path: &Path) -> Result> { Ok(get_seek_index(path)? .iter() .map(|(k, v)| (*k, v.blocks.iter().map(|b| b.size).sum::())) .collect()) } pub fn import_seek_index(input: &mut EbmlReader) -> Result> { 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, ) -> 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; while let Some(Ok(Unflat { children, item, position, })) = children.n() { match item { MatroskaTag::Timestamp(ts) => pts = ts, MatroskaTag::BlockGroup(_) => { trace!("group"); let mut children = children.unwrap(); while let Some(Ok(Unflat { children: _, item, .. })) = children.n() { match item { MatroskaTag::Block(ref block) => { debug!( "block: track={} tso={}", block.track, block.timestamp_off ); seek_index_add(seek_index, block, position.unwrap(), pts); } _ => trace!("{item:?}"), } } } MatroskaTag::SimpleBlock(block) => { trace!( "simple block: track={} tso={}", block.track, block.timestamp_off ); trace!("{pts} {}", block.timestamp_off); seek_index_add(seek_index, &block, position.unwrap(), pts); } _ => trace!("(rsc) tag ignored: {item:?}"), } } } _ => debug!("(rs) tag ignored: {item:?}"), }; } Ok(()) } fn seek_index_add( seek_index: &mut BTreeMap, block: &Block, position: u64, pts_base: u64, ) { //* I heard this helped debugging once. // { // let mut f = File::open("/home/muffin/videos/itte-yorushika.mkv").unwrap(); // f.seek(std::io::SeekFrom::Start(position.try_into().unwrap())) // .unwrap(); // let mut buf = [0u8]; // f.read_exact(&mut buf).unwrap(); // eprintln!("{}", buf[0]); // if buf[0] != 0xa0 && buf[0] != 0xa3 { // warn!("invalid position {position}") // } // } let trs = seek_index.entry(block.track).or_default(); if block.keyframe { trs.keyframes.push(trs.blocks.len()); } trs.blocks.push(BlockIndex { pts: (pts_base as i64 + block.timestamp_off as i64) as u64, source_off: position, size: block.data.len(), }); }