path: root/remuxer/src/lib.rs

/*
    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>
*/
pub mod import;
pub mod segment_extractor;
pub mod trim_writer;

use crate::{segment_extractor::SegmentExtractIter, trim_writer::TrimWriter};
use anyhow::{anyhow, Context};
use jellycommon::{BlockIndex, ItemInfo, SeekIndex, SourceTrack, SourceTrackKind};
use jellymatroska::{
    block::Block,
    read::EbmlReader,
    write::{bad_vint_length, vint_length, EbmlWriter},
    Master, MatroskaTag,
};
use log::{debug, info, trace, warn};
use std::{
    fs::File,
    io::{Seek, SeekFrom, Write},
    ops::Range,
    path::PathBuf,
    time::Instant,
};

#[derive(Debug, Clone)]
pub struct RemuxerContext {}

impl RemuxerContext {
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        Self {}
    }

    pub fn generate_into(
        &self,
        writer: impl Write + 'static,
        range: Range<usize>,
        path_base: PathBuf,
        iteminfo: ItemInfo,
        selection: Vec<usize>,
        webm: bool,
    ) -> anyhow::Result<()> {
        info!("remuxing {:?} to have tracks {selection:?}", iteminfo.title);
        let writer = TrimWriter::new(writer, range.clone());
        let mut output = EbmlWriter::new(writer, 0);

        struct ReaderC {
            info: SourceTrack,
            reader: EbmlReader,
            mapped: u64,
            index: SeekIndex,
            layouting_progress_index: usize,
        }

        let timing_cp = Instant::now();

        let mut inputs = selection
            .iter()
            .enumerate()
            .map(|(index, sel)| {
                let info = iteminfo
                    .tracks
                    .get(sel)
                    .ok_or(anyhow!("track not available"))?
                    .to_owned();
                let source_path = path_base.join(&info.path);
                let mapped = index as u64 + 1;
                info!(
                    "\t- {sel} {source_path:?} ({} => {mapped})",
                    info.track_number
                );
                info!("\t    {}", info);
                let file = File::open(&source_path).context("opening source file")?;
                let mut index =
                    File::open(source_path.with_extension(format!("si.{}", info.track_number)))
                        .context("opening seek index file")?;
                let index = bincode::decode_from_std_read::<SeekIndex, _, _>(
                    &mut index,
                    bincode::config::standard(),
                )?;
                debug!("\t    seek index: {} blocks loaded", index.blocks.len());
                let reader = EbmlReader::new(file);
                Ok(ReaderC {
                    index,
                    reader,
                    info,
                    mapped,
                    layouting_progress_index: 0,
                })
            })
            .collect::<anyhow::Result<Vec<_>>>()?;

        info!("(perf) prepare inputs: {:?}", Instant::now() - timing_cp);
        let timing_cp = Instant::now();

        output.write_tag(&MatroskaTag::Ebml(Master::Collected(vec![
            MatroskaTag::EbmlVersion(1),
            MatroskaTag::EbmlReadVersion(1),
            MatroskaTag::EbmlMaxIdLength(4),
            MatroskaTag::EbmlMaxSizeLength(8),
            MatroskaTag::DocType(if webm {
                "webm".to_string()
            } else {
                "matroska".to_string()
            }),
            MatroskaTag::DocTypeVersion(4),
            MatroskaTag::DocTypeReadVersion(2),
        ])))?;

        output.write_tag(&MatroskaTag::Segment(Master::Start))?;
        let segment_offset = output.position();

        output.write_tag(&MatroskaTag::Info(Master::Collected(vec![
            MatroskaTag::TimestampScale(1_000_000),
            MatroskaTag::Duration(iteminfo.duration * 1000.0),
            MatroskaTag::Title(iteminfo.title),
            MatroskaTag::MuxingApp("jellyremux".to_string()),
            MatroskaTag::WritingApp("jellything".to_string()),
        ])))?;
        output.write_tag(&MatroskaTag::Tags(Master::Collected(vec![])))?;

        let tracks_header = inputs
            .iter()
            .map(|rc| track_to_ebml(rc.mapped, &rc.info))
            .collect();
        output.write_tag(&MatroskaTag::Tracks(Master::Collected(tracks_header)))?;

        struct ClusterLayout {
            position: usize,
            timestamp: u64,
            source_offsets: Vec<Option<usize>>,
            blocks: Vec<(usize, BlockIndex)>,
        }

        let mut segment_layout: Vec<ClusterLayout> = {
            let mut cluster_pts = 0;
            let mut clusters = vec![];
            let mut cluster = vec![];
            let mut source_offsets = vec![None; inputs.len()];
            let mut gp = 0usize; // cluster position (in the segment)
            let mut p = 0usize; // block position (in the cluster)
            loop {
                let (track, block) = {
                    let mut best_block = BlockIndex {
                        pts: u64::MAX,
                        size: 0,
                        source_off: 0,
                    };
                    let mut best_track = 0;
                    for (i, r) in inputs.iter().enumerate() {
                        if let Some(v) = r.index.blocks.get(r.layouting_progress_index) {
                            if v.pts < best_block.pts {
                                best_block = v.to_owned();
                                best_track = i;
                            }
                        };
                    }
                    (best_track, best_block)
                };
                inputs[track].layouting_progress_index += 1;
                source_offsets[track].get_or_insert(block.source_off);
                if block.pts > cluster_pts + 1_000 {
                    let cluster_content_size = 1 + 1 // timestamp {tag, size}
                     + bad_vint_length(cluster_pts) // timestamp tag value
                     + p;
                    let cluster_size = 4 // tag length
                      + vint_length(cluster_content_size as u64) // size varint
                      + cluster_content_size;
                    clusters.push(ClusterLayout {
                        position: gp, // relative to the first cluster
                        timestamp: cluster_pts,
                        source_offsets,
                        blocks: std::mem::take(&mut cluster),
                    });

                    cluster_pts = block.pts;
                    source_offsets = vec![None; inputs.len()];
                    gp += cluster_size;
                    p = 0;
                }
                if block.pts == u64::MAX {
                    break;
                }

                let simpleblock_size = 1 + 2 + 1 // block {tracknum, pts_off, flags}
                                            // TODO does not work, if more than 127 tracks are present
                                            + block.size; // block payload
                p += 1; // simpleblock tag
                p += vint_length(simpleblock_size as u64); // simpleblock size vint
                p += simpleblock_size;

                cluster.push((track, block))
            }
            info!("segment layout computed ({} clusters)", clusters.len());
            clusters
        };
        info!("(perf) compute segment layout: {:?}", Instant::now() - timing_cp);
        let timing_cp = Instant::now();

        let max_cue_size = 4 // cues id
            + 8 // cues len
            + ( // cues content
                1 // cp id 
                + 1 // cp len
                + ( // cp content
                    1 // ctime id,
                    + 1 // ctime len
                    + 8 // ctime content uint
                    + ( // ctps
                        1 // ctp id
                        + 8 // ctp len
                        + (// ctp content
                            1 // ctrack id
                            + 1 // ctrack size
                            + 1 // ctrack content int 
                                // TODO break if inputs.len() >= 127
                            + 1 // ccp id
                            + 1 // ccp len
                            + 8 // ccp content offset
                        ) 
                    )
                ) * inputs.len()
            ) * segment_layout.len()
            + 1 // void id
            + 8; // void len

        let first_cluster_offset_predict = max_cue_size + output.position();

        // make the cluster position relative to the segment start as they should
        segment_layout.iter_mut().for_each(|e| e.position += first_cluster_offset_predict - segment_offset);

        output.write_tag(&MatroskaTag::Cues(Master::Collected(
            segment_layout
                .iter()
                .map(|cluster| {
                    MatroskaTag::CuePoint(Master::Collected(
                        Some(MatroskaTag::CueTime(cluster.timestamp))
                            .into_iter()
                            // TODO: Subtitles should not have cues for every cluster
                            .chain(inputs.iter().map(|i| {
                                MatroskaTag::CueTrackPositions(Master::Collected(vec![
                                    MatroskaTag::CueTrack(i.mapped),
                                    MatroskaTag::CueClusterPosition(cluster.position as u64),
                                ]))
                            }))
                            .collect(),
                    ))
                })
                .collect(),
        )))?;
        output.write_padding(first_cluster_offset_predict)?;
        let first_cluster_offset = output.position();
        assert_eq!(first_cluster_offset, first_cluster_offset_predict);

        let mut skip = 0;
        for (i, cluster) in segment_layout.iter().enumerate() {
            if (cluster.position + segment_offset) >= range.start {
                break;
            }
            skip = i;
        }

        if skip != 0 {
            info!("skipping {skip} clusters");
            output.seek(SeekFrom::Start((segment_layout[skip].position + segment_offset) as u64))?;
        }

        struct ReaderD<'a> {
            peek: Option<Block>,
            stream: SegmentExtractIter<'a>,
            mapped: u64,
        }

        let mut track_readers = inputs
            .iter_mut()
            .enumerate()
            .map(|(i, inp)| {
                inp.reader
                    .seek(
                        segment_layout[skip].source_offsets[i].unwrap(), // TODO will crash if there is a "hole"
                        MatroskaTag::Cluster(Master::Start),
                    )
                    .context("seeking in input")?;
                let mut stream = SegmentExtractIter::new(&mut inp.reader, inp.info.track_number);

                Ok(ReaderD {
                    mapped: inp.mapped,
                    peek: Some(stream.next()?),
                    stream,
                })
            })
            .collect::<anyhow::Result<Vec<_>>>()?;

        info!("(perf) seek inputs: {:?}", Instant::now() - timing_cp);

        for (cluster_index, cluster) in segment_layout.into_iter().enumerate().skip(skip) {
            debug!(
                "writing cluster {cluster_index} (pts_base={}) with {} blocks",
                cluster.timestamp,
                cluster.blocks.len()
            );
            {
                let cue_error =
                    cluster.position as i64 - (output.position() - segment_offset) as i64;
                if cue_error != 0 {
                    warn!("calculation was {} bytes off", cue_error);
                }
            }

            let mut cluster_blocks = vec![MatroskaTag::Timestamp(cluster.timestamp)];
            for (block_track, index_block) in cluster.blocks {
                let track_reader = &mut track_readers[block_track];
                let mut block = track_reader
                    .peek
                    .replace(track_reader.stream.next()?)
                    .expect("source file too short");

                assert_eq!(index_block.size, block.data.len(), "seek index is wrong");

                block.track = track_reader.mapped;
                block.timestamp_off = (index_block.pts - cluster.timestamp).try_into().unwrap();
                trace!("n={} tso={}", block.track, block.timestamp_off);
                let buf = block.dump();
                cluster_blocks.push(MatroskaTag::SimpleBlock(buf))
            }
            output.write_tag(&MatroskaTag::Cluster(Master::Collected(cluster_blocks)))?;
        }
        output.write_tag(&MatroskaTag::Segment(Master::End))?;
        Ok(())
    }
}

pub fn track_to_ebml(number: u64, track: &SourceTrack) -> MatroskaTag {
    let mut els = vec![
        MatroskaTag::TrackNumber(number),
        MatroskaTag::TrackUID(number),
        MatroskaTag::FlagLacing(0),
        MatroskaTag::Language(track.language.clone()),
        MatroskaTag::CodecID(track.codec.clone()),
    ];
    if let Some(d) = &track.default_duration {
        els.push(MatroskaTag::DefaultDuration(*d));
    }
    match track.kind {
        SourceTrackKind::Video {
            width,
            height,
            fps: _,
        } => {
            els.push(MatroskaTag::TrackType(1));
            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()),
            ])));
            els.push(MatroskaTag::BitDepth(bit_depth.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))
}