stream info for transcoding

1 files changed, 164 insertions, 0 deletions

diff --git a/stream/src/stream_info.rs b/stream/src/stream_info.rs
new file mode 100644
index 0000000..9d3d741
--- /dev/null
+++ b/stream/src/stream_info.rs
@@ -0,0 +1,164 @@
+use anyhow::Result;
+use ebml_struct::matroska::TrackEntry;
+use jellybase::{
+    common::stream::{
+        StreamContainer, StreamFormatInfo, StreamInfo, StreamSegmentInfo, StreamTrackInfo,
+        TrackKind,
+    },
+    CONF,
+};
+use jellyremuxer::metadata::{matroska_metadata, MatroskaMetadata};
+use std::{path::PathBuf, sync::Arc};
+use tokio::{
+    io::{AsyncWriteExt, DuplexStream},
+    spawn,
+    task::spawn_blocking,
+};
+
+use crate::SMediaInfo;
+
+async fn async_matroska_metadata(path: PathBuf) -> Result<Arc<MatroskaMetadata>> {
+    Ok(spawn_blocking(move || matroska_metadata(&path)).await??)
+}
+
+pub(crate) struct InternalStreamInfo {
+    pub paths: Vec<PathBuf>,
+    pub metadata: Vec<Arc<MatroskaMetadata>>,
+    pub track_to_file: Vec<(usize, u64)>,
+}
+
+pub(crate) async fn stream_info(info: Arc<SMediaInfo>) -> Result<(InternalStreamInfo, StreamInfo)> {
+    let mut metadata = Vec::new();
+    let mut paths = Vec::new();
+    for path in &info.files {
+        metadata.push(async_matroska_metadata(path.clone()).await?);
+        paths.push(path.clone());
+    }
+    let mut tracks = Vec::new();
+    let mut track_to_file = Vec::new();
+
+    for (i, m) in metadata.iter().enumerate() {
+        if let Some(t) = &m.tracks {
+            for t in &t.entries {
+                tracks.push(StreamTrackInfo {
+                    name: None,
+                    kind: match t.track_type {
+                        1 => TrackKind::Video,
+                        2 => TrackKind::Audio,
+                        17 => TrackKind::Subtitle,
+                        _ => todo!(),
+                    },
+                    formats: stream_formats(t),
+                });
+                track_to_file.push((i, t.track_number));
+            }
+        }
+    }
+
+    let segment = StreamSegmentInfo {
+        name: None,
+        duration: metadata[0]
+            .info
+            .as_ref()
+            .unwrap()
+            .duration
+            .unwrap_or_default(),
+        tracks,
+    };
+    Ok((
+        InternalStreamInfo {
+            metadata,
+            paths,
+            track_to_file,
+        },
+        StreamInfo {
+            name: info.info.title.clone(),
+            segments: vec![segment],
+        },
+    ))
+}
+
+fn stream_formats(t: &TrackEntry) -> Vec<StreamFormatInfo> {
+    let mut formats = Vec::new();
+    formats.push(StreamFormatInfo {
+        codec: t.codec_id.to_string(),
+        remux: true,
+        bitrate: 2_000_000., // TODO
+        containers: containers_by_codec(&t.codec_id),
+        bit_depth: t.audio.as_ref().and_then(|a| a.bit_depth.map(|e| e as u8)),
+        samplerate: t.audio.as_ref().map(|a| a.sampling_frequency),
+        channels: t.audio.as_ref().map(|a| a.channels as usize),
+        width: t.video.as_ref().map(|v| v.pixel_width),
+        height: t.video.as_ref().map(|v| v.pixel_height),
+        ..Default::default()
+    });
+
+    match t.track_type {
+        1 => {
+            let sw = t.video.as_ref().unwrap().pixel_width;
+            let sh = t.video.as_ref().unwrap().pixel_height;
+            for (w, br) in [(3840, 8e6), (1920, 5e6), (1280, 3e6), (640, 1e6)] {
+                if w > sw {
+                    continue;
+                }
+                let h = (w * sh) / sw;
+                for (cid, enable) in [
+                    ("V_AV1", CONF.encoders.av1.is_some()),
+                    ("V_VP8", CONF.encoders.vp8.is_some()),
+                    ("V_VP9", CONF.encoders.vp9.is_some()),
+                    ("V_AVC", CONF.encoders.avc.is_some()),
+                    ("V_HEVC", CONF.encoders.hevc.is_some()),
+                ] {
+                    if enable {
+                        formats.push(StreamFormatInfo {
+                            codec: cid.to_string(),
+                            bitrate: br,
+                            remux: false,
+                            containers: containers_by_codec(cid),
+                            width: Some(w),
+                            height: Some(h),
+                            samplerate: None,
+                            channels: None,
+                            bit_depth: None,
+                        });
+                    }
+                }
+            }
+        }
+        2 => {
+            for br in [256e3, 128e3, 64e3] {
+                formats.push(StreamFormatInfo {
+                    codec: "A_OPUS".to_string(),
+                    bitrate: br,
+                    remux: false,
+                    containers: containers_by_codec("A_OPUS"),
+                    width: None,
+                    height: None,
+                    samplerate: Some(48e3),
+                    channels: Some(2),
+                    bit_depth: Some(32),
+                });
+            }
+        }
+        17 => {}
+        _ => {}
+    }
+
+    formats
+}
+
+fn containers_by_codec(codec: &str) -> Vec<StreamContainer> {
+    use StreamContainer::*;
+    match codec {
+        "V_VP8" | "V_VP9" | "V_AV1" | "A_OPUS" | "A_VORBIS" => vec![Matroska, WebM],
+        "V_AVC" | "A_AAC" => vec![Matroska, MPEG4],
+        "S_TEXT/UTF8" | "S_TEXT/WEBVTT" => vec![Matroska, WebVTT, WebM, JVTT],
+        _ => vec![Matroska],
+    }
+}
+
+pub(crate) async fn write_stream_info(info: Arc<SMediaInfo>, mut b: DuplexStream) -> Result<()> {
+    let (_, info) = stream_info(info).await?;
+    spawn(async move { b.write_all(&serde_json::to_vec(&info)?).await });
+    Ok(())
+}