use anyhow::{anyhow, Context}; use ffmpeg_next::{ format::{input, stream::Disposition}, media::Type, Stream, }; use lazy_static::lazy_static; use log::debug; use regex::Regex; use tempfile::tempdir; use tokio::{fs, process::Command}; use std::path::Path; use crate::config::CaptureConfig; lazy_static! { static ref SUBTITLE_FORBID_REGEX: Regex = Regex::new("(?i)sign|song").unwrap(); } pub struct MediaInfo { pub duration_secs: f64, // The index among the subtitle streams, not among the streams in general pub subtitle_stream_index: Option, // The index among the audio streams, not among the streams in general pub audio_stream_index: Option, } fn indexed_streams( ctx: &ffmpeg_next::format::context::common::Context, stream_type: Type, ) -> impl Iterator)> { ctx.streams() .filter(move |stream| { ffmpeg_next::codec::context::Context::from_parameters(stream.parameters()) .map(|c| c.medium()) == Ok(stream_type) }) .enumerate() } pub fn get_media_info>( source: &P, capture_config: &CaptureConfig, ) -> anyhow::Result { let ctx = input(source).context("Failed to load media file")?; let duration_secs = ctx.duration() as f64 / f64::from(ffmpeg_next::ffi::AV_TIME_BASE); debug!("{:?}", ctx.metadata()); let subtitle_stream_index = capture_config.subtitle_language.as_ref().and_then(|lang| { indexed_streams(&ctx, Type::Subtitle) .filter(|(_, stream)| { let metadata = stream.metadata(); if metadata.get("language") != Some(lang) { return false; } if metadata .get("title") .map(|t| SUBTITLE_FORBID_REGEX.is_match(t)) == Some(true) { return false; } true }) .min_by_key(|(_, stream)| stream.disposition().contains(Disposition::FORCED)) .map(|(idx, _)| idx) }); let audio_stream_index = capture_config.audio_language.as_ref().and_then(|lang| { indexed_streams(&ctx, Type::Audio) .find(|(_, stream)| stream.metadata().get("language") == Some(lang)) .map(|(idx, _)| idx) }); Ok(MediaInfo { duration_secs, subtitle_stream_index, audio_stream_index, }) } async fn take_ffmpeg_capture( source: &P, output_ext: &str, apply_args: F, ) -> anyhow::Result> where P: AsRef, F: FnOnce(&mut Command, &Path, &Path), { let input_ext = source.as_ref().extension().and_then(|s| s.to_str()); if input_ext.map(|e| e.chars().all(|c| c.is_ascii_alphanumeric())) != Some(true) { return Err(anyhow!( "Media file {} had unexpected file extension", source.as_ref().display() )); } let tmp_dir = tempdir() .context("Failed to create temporary directory for ffmpeg input and output files")?; let link_path = tmp_dir.path().join(format!("in.{}", input_ext.unwrap())); fs::symlink(source, &link_path) .await .context("Failed to create symlink for input file")?; let dest_path = tmp_dir.path().join(format!("out.{}", output_ext)); let mut cmd = Command::new("ffmpeg"); apply_args(&mut cmd, &link_path, &dest_path); let status = cmd .status() .await .with_context(|| format!("Error running ffmpeg child process {:?}", cmd))?; if !status.success() { match status.code() { Some(code) => return Err(anyhow!("ffmpeg exited with status code {code}")), None => return Err(anyhow!("ffmpeg terminated by signal")), } } fs::read(&dest_path) .await .with_context(|| format!("Failed to read ffmpeg output file {}", dest_path.display())) } pub async fn take_screencap>( source: &P, timestamp_secs: f64, subtitle_stream_index: Option, ) -> anyhow::Result> { take_ffmpeg_capture(source, "png", |cmd, in_path, out_path| { cmd.arg("-ss") .arg(format!("{:.2}", timestamp_secs)) .arg("-copyts") .arg("-i") .arg(in_path); if let Some(idx) = subtitle_stream_index { cmd.arg("-filter_complex").arg(format!( "[0:v]subtitles={}:si={}", in_path.to_string_lossy(), idx )); } cmd.args(["-vframes", "1"]) .args(["-loglevel", "quiet"]) .arg("-y") .arg(out_path); }) .await } pub async fn take_audio_clip>( source: &P, timestamp_secs: f64, duration_secs: f64, audio_stream_index: Option, ) -> anyhow::Result> { take_ffmpeg_capture(source, "mp3", |cmd, in_path, out_path| { cmd.arg("-ss") .arg(format!("{:.2}", timestamp_secs)) .arg("-t") .arg(format!("{:.2}", duration_secs)) .arg("-i") .arg(in_path); if let Some(idx) = audio_stream_index { cmd.arg("-map").arg(format!("0:a:{}", idx)); } cmd.args(["-loglevel", "quiet"]).arg("-y").arg(out_path); }) .await }