path: root/remuxer/src/import/mod.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>
*/
use anyhow::{anyhow, bail, Result};
use jellycommon::{BlockIndex, LocalTrack, SeekIndex, SourceTrack, SourceTrackKind};
use jellymatroska::{
    block::Block,
    matroska::MatroskaTag,
    read::EbmlReader,
    unflatten::{IterWithPos, Unflat, Unflatten},
};
use log::{debug, error, info, trace, warn};
use std::{collections::BTreeMap, path::PathBuf};

pub fn import_read(
    path: &PathBuf,
    input: &mut EbmlReader,
) -> Result<(Vec<SourceTrack>, Vec<LocalTrack>, BTreeMap<u64, SeekIndex>)> {
    let mut iteminfo = Vec::new();
    let mut private = Vec::new();
    let mut seek_index = BTreeMap::new();
    while let Some(item) = input.next() {
        let item = match item {
            Ok(item) => item,
            Err(e) => {
                warn!("{e}");
                break;
            }
        };
        match item {
            MatroskaTag::Ebml(_) => {
                let mut iter = Unflatten::new_with_end(input, item);
                while let Some(Ok(Unflat {
                    children: _, item, ..
                })) = iter.n()
                {
                    match item {
                        MatroskaTag::DocType(t) => {
                            if !matches!(t.as_str(), "matroska" | "webm") {
                                error!("file is neither matroska nor webm but {:?}", t)
                            }
                        }
                        _ => debug!("(re) tag ignored: {item:?}"),
                    }
                }
            }
            MatroskaTag::Segment(_) => {
                info!("segment start");
                let mut children = Unflatten::new_with_end(input, item);
                import_read_segment(
                    path,
                    &mut children,
                    &mut iteminfo,
                    &mut private,
                    &mut seek_index,
                )?;
                info!("segment end");
            }
            _ => debug!("(r) tag ignored: {item:?}"),
        }
    }

    Ok((iteminfo, private, seek_index))
}

fn import_read_segment(
    path: &PathBuf,
    segment: &mut Unflatten,
    iteminfo: &mut Vec<SourceTrack>,
    private: &mut Vec<LocalTrack>,
    seek_index: &mut BTreeMap<u64, SeekIndex>,
) -> Result<Option<f64>> {
    let (mut timestamp_scale, mut duration) = (None, None);

    while let Some(Ok(Unflat { children, item, .. })) = segment.n() {
        match item {
            MatroskaTag::SeekHead(_) => {}
            MatroskaTag::Info(_) => {
                let mut children = children.unwrap();
                while let Some(Ok(Unflat {
                    children: _, item, ..
                })) = children.n()
                {
                    match item {
                        MatroskaTag::TimestampScale(v) => timestamp_scale = Some(v),
                        MatroskaTag::Duration(v) => duration = Some(v),
                        _ => debug!("(rsi) tag ignored: {item:?}"),
                    }
                }
            }
            MatroskaTag::Tags(_) => {}
            MatroskaTag::Cues(_) => {}
            MatroskaTag::Chapters(_) => {}
            MatroskaTag::Tracks(_) => {
                let mut children = children.unwrap();
                while let Some(Ok(Unflat { children, item, .. })) = children.n() {
                    match item {
                        MatroskaTag::TrackEntry(_) => {
                            let mut children = children.unwrap();
                            let (
                                mut index,
                                mut language,
                                mut codec,
                                mut kind,
                                mut sample_rate,
                                mut channels,
                                mut width,
                                mut height,
                                mut name,
                                mut fps,
                                mut bit_depth,
                                mut codec_private,
                                mut default_duration,
                            ) = (
                                None, None, None, None, None, None, None, None, None, None, None,
                                None, None,
                            );
                            while let Some(Ok(Unflat { children, item, .. })) = children.n() {
                                match item {
                                    MatroskaTag::CodecID(b) => codec = Some(b),
                                    MatroskaTag::Language(v) => language = Some(v),
                                    MatroskaTag::TrackNumber(v) => index = Some(v),
                                    MatroskaTag::TrackType(v) => kind = Some(v),
                                    MatroskaTag::Name(v) => name = Some(v),
                                    MatroskaTag::CodecPrivate(v) => codec_private = Some(v),
                                    MatroskaTag::DefaultDuration(v) => default_duration = Some(v),
                                    MatroskaTag::Audio(_) => {
                                        let mut children = children.unwrap();
                                        while let Some(Ok(Unflat { item, .. })) = children.n() {
                                            match item {
                                                MatroskaTag::Channels(v) => {
                                                    channels = Some(v as usize)
                                                }
                                                MatroskaTag::SamplingFrequency(v) => {
                                                    sample_rate = Some(v)
                                                }
                                                MatroskaTag::BitDepth(v) => bit_depth = Some(v),
                                                _ => (),
                                            }
                                        }
                                    }
                                    MatroskaTag::Video(_) => {
                                        let mut children = children.unwrap();
                                        while let Some(Ok(Unflat { item, .. })) = children.n() {
                                            match item {
                                                MatroskaTag::PixelWidth(v) => width = Some(v),
                                                MatroskaTag::PixelHeight(v) => height = Some(v),
                                                MatroskaTag::FrameRate(v) => fps = Some(v),
                                                _ => (),
                                            }
                                        }
                                    }
                                    _ => (),
                                }
                            }
                            let track_index = index.unwrap();
                            let kind = match kind.ok_or(anyhow!("track type required"))? {
                                1 => SourceTrackKind::Video {
                                    fps: fps.unwrap_or(0.0), // TODO
                                    width: width.unwrap(),
                                    height: height.unwrap(),
                                },
                                2 => SourceTrackKind::Audio {
                                    bit_depth: bit_depth.unwrap_or(0) as usize, // TODO
                                    channels: channels.unwrap(),
                                    sample_rate: sample_rate.unwrap(),
                                },
                                17 => SourceTrackKind::Subtitles,
                                _ => bail!("invalid track type"),
                            };
                            iteminfo.push(SourceTrack {
                                default_duration,
                                name: name.unwrap_or_else(|| "unnamed".to_string()),
                                codec: codec.unwrap(),
                                language: language.unwrap_or_else(|| "none".to_string()),
                                kind,
                            });
                            private.push(LocalTrack {
                                track: track_index as usize,
                                path: path.to_owned(),
                                codec_private,
                            })
                        }
                        _ => debug!("(rst) tag ignored: {item:?}"),
                    }
                }
            }
            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(_) => {
                                debug!("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
                                                .entry(block.track)
                                                .or_insert(SeekIndex { blocks: vec![] })
                                                .blocks
                                                .push(BlockIndex {
                                                    pts: pts + block.timestamp_off as u64,
                                                    source_off: position,
                                                    size: block.data.len(),
                                                });
                                        }
                                        _ => trace!("{item:?}"),
                                    }
                                }
                            }
                            MatroskaTag::SimpleBlock(buf) => {
                                let block = Block::parse(&buf)?;
                                debug!(
                                    "simple block: track={} tso={}",
                                    block.track, block.timestamp_off
                                );
                                debug!("{pts} {}", block.timestamp_off);
                                seek_index
                                    .entry(block.track)
                                    .or_insert(SeekIndex { blocks: vec![] })
                                    .blocks
                                    .push(BlockIndex {
                                        pts: (pts as i64 + block.timestamp_off as i64) as u64,
                                        source_off: position,
                                        size: block.data.len(),
                                    });
                            }
                            _ => debug!("(rsc) tag ignored: {item:?}"),
                        }
                    } else {
                        break;
                    }
                    position = children.position();
                }
            }
            _ => debug!("(rs) tag ignored: {item:?}"),
        };
    }

    Ok(if let Some(duration) = duration {
        Some((duration * timestamp_scale.unwrap_or(1_000_000) as f64) / 1_000_000_000_f64)
    } else {
        None
    })
}