diff options
Diffstat (limited to 'remuxer/src/lib.rs')
-rw-r--r-- | remuxer/src/lib.rs | 249 |
1 files changed, 180 insertions, 69 deletions
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 { |