shrivel 0.4.0

Cross-platform FFmpeg orchestrator that transcodes videos to HEVC or AV1
//! Per-file conversion decisions and FFmpeg argument construction.

use std::ffi::OsString;
use std::path::Path;

use crate::cli::{AudioCodec, Codec};
use crate::encoder::Encoder;

/// A frame-rate difference at or below this value is not worth resampling.
const FPS_TOLERANCE: f64 = 0.5;

/// Containers that need the `hvc1` tag for HEVC to play on Apple devices.
const HVC1_CONTAINERS: [&str; 3] = ["mp4", "m4v", "mov"];

/// Returns true when the source frame rate differs enough from the target
/// to justify resampling.
pub fn fps_change_needed(source: f64, target: f64) -> bool {
    (source - target).abs() > FPS_TOLERANCE
}

/// Calculates a dynamic CQ value based on the source bitrate.
pub fn calculate_cq(base_cq: u8, bitrate_kbps: u32) -> u8 {
    let offset = if bitrate_kbps <= 2000 {
        4
    } else if bitrate_kbps <= 4000 {
        2
    } else if bitrate_kbps <= 8000 {
        0
    } else if bitrate_kbps <= 12000 {
        -2
    } else {
        -4
    };
    (base_cq as i16 + offset).clamp(0, 51) as u8
}

/// Everything needed to build one FFmpeg invocation.
pub struct Plan<'a> {
    pub codec: Codec,
    pub encoder: Encoder,
    pub input: &'a Path,
    pub output: &'a Path,
    pub cq: u8,
    pub audio_codec: AudioCodec,
    pub no_audio: bool,
    /// Video filters, applied in order.
    pub filters: Vec<String>,
    /// Force constant frame rate output (used when resampling).
    pub force_cfr: bool,
}

impl Plan<'_> {
    /// Builds the full FFmpeg argument list (without the program name).
    pub fn ffmpeg_args(&self) -> Vec<OsString> {
        let mut args: Vec<OsString> = Vec::new();
        let mut push = |s: &str| args.push(OsString::from(s));

        for a in [
            "-hide_banner",
            "-loglevel",
            "error",
            "-nostdin",
            "-nostats",
            "-progress",
            "pipe:2",
        ] {
            push(a);
        }
        for a in self.encoder.input_args() {
            push(a);
        }
        push("-i");
        args.push(self.input.as_os_str().to_owned());

        let mut push = |s: &str| args.push(OsString::from(s));
        for a in ["-map", "0:v:0", "-map", "0:s?"] {
            push(a);
        }
        if !self.no_audio {
            for a in ["-map", "0:a?"] {
                push(a);
            }
        }
        for a in self.encoder.video_args(self.codec, self.cq) {
            push(&a);
        }
        if !self.filters.is_empty() {
            push("-vf");
            push(&self.filters.join(","));
        }
        if self.codec == Codec::Hevc && self.wants_hvc1_tag() {
            push("-tag:v");
            push("hvc1");
        }
        if self.no_audio {
            push("-an");
        } else {
            match self.audio_codec {
                AudioCodec::Copy => {
                    for a in ["-c:a", "copy"] {
                        push(a);
                    }
                }
                AudioCodec::Mp3 => {
                    self.audio_encode_args(&mut push, "libmp3lame", audio_bitrate(self.cq, 320, 96))
                }
                AudioCodec::Aac => {
                    self.audio_encode_args(&mut push, "aac", audio_bitrate(self.cq, 256, 64))
                }
                AudioCodec::Opus => {
                    self.audio_encode_args(&mut push, "libopus", audio_bitrate(self.cq, 192, 48))
                }
            }
        }
        for a in ["-c:s", "copy"] {
            push(a);
        }
        if self.force_cfr {
            push("-fps_mode");
            push("cfr");
        }
        push("-y");
        args.push(self.output.as_os_str().to_owned());
        args
    }

    fn audio_encode_args(&self, push: &mut impl FnMut(&str), encoder: &str, bitrate: u16) {
        push("-c:a");
        push(encoder);
        push("-b:a");
        // FFmpeg accepts bitrate values as strings; format locally then append below.
        push(&format!("{bitrate}k"));
    }

    fn wants_hvc1_tag(&self) -> bool {
        self.output
            .extension()
            .and_then(|e| e.to_str())
            .is_some_and(|ext| HVC1_CONTAINERS.iter().any(|c| c.eq_ignore_ascii_case(ext)))
    }
}

/// Map the normalized quality scale to a practical codec-specific bitrate.
/// CQ 0 uses the high end and CQ 51 the low end.
fn audio_bitrate(cq: u8, high: u16, low: u16) -> u16 {
    high - (((high - low) as u32 * cq.min(51) as u32 + 25) / 51) as u16
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;

    #[test]
    fn fps_tolerance() {
        assert!(!fps_change_needed(30.0, 30.0));
        assert!(!fps_change_needed(29.97, 30.0));
        assert!(fps_change_needed(60.0, 30.0));
        // An unknown source rate (0) always triggers a change.
        assert!(fps_change_needed(0.0, 30.0));
    }

    fn plan<'a>(input: &'a PathBuf, output: &'a PathBuf) -> Plan<'a> {
        Plan {
            codec: Codec::Hevc,
            encoder: Encoder::Software,
            input,
            output,
            cq: 30,
            audio_codec: AudioCodec::Copy,
            no_audio: false,
            filters: vec![],
            force_cfr: false,
        }
    }

    fn strings(args: &[OsString]) -> Vec<String> {
        args.iter()
            .map(|a| a.to_string_lossy().into_owned())
            .collect()
    }

    #[test]
    fn basic_args_end_with_overwrite_and_output() {
        let (i, o) = (PathBuf::from("in.mp4"), PathBuf::from("out.mp4"));
        let args = strings(&plan(&i, &o).ffmpeg_args());
        assert_eq!(&args[args.len() - 2..], ["-y", "out.mp4"]);
        assert!(args.windows(2).any(|w| w == ["-i", "in.mp4"]));
        assert!(args.windows(2).any(|w| w == ["-tag:v", "hvc1"]));
        assert!(!args.contains(&"-vf".to_string()));
        assert!(!args.contains(&"-fps_mode".to_string()));
    }

    #[test]
    fn default_audio_is_mapped_and_copied() {
        let (i, o) = (PathBuf::from("in.mp4"), PathBuf::from("out.mp4"));
        let args = strings(&plan(&i, &o).ffmpeg_args());
        assert!(args.windows(2).any(|w| w == ["-map", "0:a?"]));
        assert!(args.windows(2).any(|w| w == ["-c:a", "copy"]));
        assert!(!args.contains(&"-an".to_string()));
    }

    #[test]
    fn subtitles_are_mapped_when_present_and_copied() {
        let (i, o) = (PathBuf::from("in.mkv"), PathBuf::from("out.mkv"));
        let args = strings(&plan(&i, &o).ffmpeg_args());
        assert!(args.windows(2).any(|w| w == ["-map", "0:s?"]));
        assert!(args.windows(2).any(|w| w == ["-c:s", "copy"]));
    }

    #[test]
    fn audio_codecs_select_the_expected_encoder_and_bitrate() {
        let (i, o) = (PathBuf::from("in.mp4"), PathBuf::from("out.mp4"));
        for (codec, encoder) in [
            (AudioCodec::Mp3, "libmp3lame"),
            (AudioCodec::Aac, "aac"),
            (AudioCodec::Opus, "libopus"),
        ] {
            let mut p = plan(&i, &o);
            p.audio_codec = codec;
            let args = strings(&p.ffmpeg_args());
            assert!(args.windows(2).any(|w| w == ["-c:a", encoder]));
            let bitrate = args.iter().position(|arg| arg == "-b:a").unwrap();
            assert!(args[bitrate + 1].ends_with('k'));
        }
    }

    #[test]
    fn audio_bitrate_uses_cq_endpoints_and_decreases_with_cq() {
        for (high, low) in [(320, 96), (256, 64), (192, 48)] {
            assert_eq!(audio_bitrate(0, high, low), high);
            assert_eq!(audio_bitrate(51, high, low), low);
            assert!(audio_bitrate(40, high, low) < audio_bitrate(10, high, low));
            assert!(audio_bitrate(26, high, low) > 0);
        }

        // A CQ adjusted by the same dynamic CQ calculation used by the
        // conversion flow must drive the audio bitrate too.
        let (i, o) = (PathBuf::from("in.mp4"), PathBuf::from("out.mp4"));
        let mut p = plan(&i, &o);
        p.cq = calculate_cq(30, 2_000);
        p.audio_codec = AudioCodec::Opus;
        let args = strings(&p.ffmpeg_args());
        let bitrate = args.iter().position(|arg| arg == "-b:a").unwrap();
        assert_eq!(
            args[bitrate + 1],
            format!("{}k", audio_bitrate(p.cq, 192, 48))
        );
    }

    #[test]
    fn no_audio_disables_mapping_even_with_selected_audio_codec() {
        let (i, o) = (PathBuf::from("in.mp4"), PathBuf::from("out.mp4"));
        let mut p = plan(&i, &o);
        p.audio_codec = AudioCodec::Opus;
        p.no_audio = true;
        let args = strings(&p.ffmpeg_args());
        assert!(!args.windows(2).any(|w| w == ["-map", "0:a?"]));
        assert!(args.contains(&"-an".to_string()));
        assert!(!args.contains(&"-c:a".to_string()));
        assert!(!args.contains(&"-b:a".to_string()));
    }

    #[test]
    fn filters_and_cfr_are_included() {
        let (i, o) = (PathBuf::from("in.mp4"), PathBuf::from("out.mp4"));
        let mut p = plan(&i, &o);
        p.filters = vec!["hqdn3d=3:3:8:8".into(), "fps=30".into()];
        p.force_cfr = true;
        let args = strings(&p.ffmpeg_args());
        assert!(
            args.windows(2)
                .any(|w| w == ["-vf", "hqdn3d=3:3:8:8,fps=30"])
        );
        let cfr = args.iter().position(|a| a == "-fps_mode").unwrap();
        assert_eq!(args[cfr + 1], "cfr");
        // Output options must come before the output path.
        assert!(cfr < args.len() - 1);
    }

    #[test]
    fn hvc1_tag_is_skipped_for_matroska() {
        let (i, o) = (PathBuf::from("in.mkv"), PathBuf::from("out.mkv"));
        let args = strings(&plan(&i, &o).ffmpeg_args());
        assert!(!args.contains(&"-tag:v".to_string()));
    }
}