screencap-bot/src/media.rs

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<usize>,
    // The index among the audio streams, not among the streams in general
    pub audio_stream_index: Option<usize>,
}

fn indexed_streams(
    ctx: &ffmpeg_next::format::context::common::Context,
    stream_type: Type,
) -> impl Iterator<Item = (usize, Stream<'_>)> {
    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<P: AsRef<Path>>(
    source: &P,
    capture_config: &CaptureConfig,
) -> anyhow::Result<MediaInfo> {
    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<P, F>(
    source: &P,
    output_ext: &str,
    apply_args: F,
) -> anyhow::Result<Vec<u8>>
where
    P: AsRef<Path>,
    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<P: AsRef<Path>>(
    source: &P,
    timestamp_secs: f64,
    subtitle_stream_index: Option<usize>,
) -> anyhow::Result<Vec<u8>> {
    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<P: AsRef<Path>>(
    source: &P,
    timestamp_secs: f64,
    duration_secs: f64,
    audio_stream_index: Option<usize>,
) -> anyhow::Result<Vec<u8>> {
    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
}