path: root/remuxer/src/fragment.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) 2025 metamuffin <metamuffin.org>
*/

use crate::{
    metadata::matroska_metadata, seek_index::get_seek_index, segment_extractor::SegmentExtractIter,
};
use anyhow::{anyhow, Context, Result};
use ebml_struct::{
    matroska::{BlockGroup, Cluster, Ebml, Info, Segment, Tracks},
    write::TagWrite,
    Block,
};
use jellybase::common::{LocalTrack, Node, SourceTrackKind};
use jellymatroska::{read::EbmlReader, 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: &Node,
    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_frag = (FRAGMENT_LENGTH / average_kf_interval).ceil() as usize;
    debug!("average keyframe interval: {average_kf_interval}");
    debug!(" => keyframes per frag {kf_per_frag}");

    let n_frags = n_kf.div_ceil(kf_per_frag);
    Ok((0..n_frags)
        .map(|i| {
            let start = index.blocks[if force_kf {
                i * kf_per_frag
            } else {
                index.keyframes[i * kf_per_frag]
            }]
            .pts as f64
                / 1000.;
            let end = if force_kf {
                let n = (i + 1) * kf_per_frag;
                if n >= index.blocks.len() {
                    None
                } else {
                    Some(n)
                }
            } else {
                index.keyframes.get((i + 1) * kf_per_frag).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: &Node,
    local_track: &LocalTrack,
    track: usize,
    webm: bool,
    n: usize,
) -> anyhow::Result<()> {
    info!("writing fragment {n} of {:?} (track {track})", item.title);

    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_frag = (FRAGMENT_LENGTH / average_kf_interval).ceil() as usize;
    debug!("average keyframe interval: {average_kf_interval}");
    debug!(" => keyframes per frag {kf_per_frag}");

    let (start_block_index, end_block_index) = if force_kf {
        (n * kf_per_frag, (n + 1) * kf_per_frag)
    } else {
        (
            *index
                .keyframes
                .get(n * kf_per_frag)
                .ok_or(anyhow!("fragment index out of range"))?,
            *index
                .keyframes
                .get((n + 1) * kf_per_frag)
                .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);

    let input_metadata = (*matroska_metadata(&local_track.path)?).clone().unwrap();

    reader.seek(start_block.source_off, MatroskaTag::Cluster(Master::Start))?;
    let mut reader = SegmentExtractIter::new(&mut reader, local_track.track as u64);

    let mut cluster = Cluster::default();
    cluster.timestamp = start_block.pts;
    for i in start_block_index..end_block_index {
        let index_block = &index.blocks[i];
        let (block, duration) = reader.next_block()?;

        let mut block = Block {
            data: block.data,
            discardable: block.discardable,
            invisible: block.invisible,
            keyframe: block.keyframe,
            lacing: block.lacing.map(|l| match l {
                jellymatroska::block::LacingType::Xiph => ebml_struct::LacingType::Xiph,
                jellymatroska::block::LacingType::FixedSize => ebml_struct::LacingType::FixedSize,
                jellymatroska::block::LacingType::Ebml => ebml_struct::LacingType::Ebml,
            }),
            timestamp_off: block.timestamp_off,
            track: block.track,
        };

        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 {
            cluster.block_groups.push(BlockGroup {
                block_duration: Some(duration),
                block,
                ..Default::default()
            })
        } else {
            cluster.simple_blocks.push(block)
        }
    }

    let mut input_track = input_metadata
        .tracks
        .unwrap()
        .entries
        .into_iter()
        .find(|t| t.track_number == local_track.track as u64)
        .unwrap();

    input_track.track_number = 1;
    if webm {
        if let Some(v) = &mut input_track.video {
            v.colour = None;
        }
    }

    let mut output = BufWriter::new(writer);

    Ebml {
        ebml_version: 1,
        ebml_read_version: 1,
        ebml_max_id_length: 4,
        ebml_max_size_length: 8,
        doc_type: if webm {
            "webm".to_string()
        } else {
            "matroska".to_string()
        },
        doc_type_version: 4,
        doc_type_read_version: 2,
        doc_type_extensions: vec![],
    }
    .write(&mut output)?;

    Segment {
        info: Info {
            timestamp_scale: 1_000_000,
            duration: Some((last_block_pts - start_block.pts) as f64),
            title: Some(format!(
                "{}: {info}",
                item.title.clone().unwrap_or_default()
            )),
            muxing_app: "ebml-struct".to_owned(),
            writing_app: "jellything".to_owned(),
            ..Default::default()
        },
        tracks: Some(Tracks {
            entries: vec![input_track],
        }),
        clusters: vec![cluster],
        ..Default::default()
    }
    .write(&mut output)?;

    Ok(())
}