shrivel 0.2.1

Cross-platform FFmpeg orchestrator that transcodes videos to HEVC or AV1
//! FFmpeg encoder backends and codec-specific arguments.

use std::process::{Command, Stdio};

use crate::cli::{Codec, EncoderChoice};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Encoder {
    Nvenc,
    VideoToolbox,
    Qsv,
    Amf,
    Software,
}

impl Encoder {
    const HEVC_HARDWARE: [Encoder; 4] = [Self::Nvenc, Self::VideoToolbox, Self::Qsv, Self::Amf];
    const AV1_HARDWARE: [Encoder; 3] = [Self::Nvenc, Self::Qsv, Self::Amf];

    pub fn from_choice(choice: EncoderChoice, codec: Codec) -> Option<Self> {
        match choice {
            EncoderChoice::Auto => None,
            EncoderChoice::Nvenc => Some(Self::Nvenc),
            EncoderChoice::Qsv => Some(Self::Qsv),
            EncoderChoice::Amf => Some(Self::Amf),
            EncoderChoice::Videotoolbox => Some(Self::VideoToolbox),
            EncoderChoice::Software => Some(Self::Software),
        }
        .filter(|encoder| encoder.supports(codec))
    }

    pub fn supports(self, codec: Codec) -> bool {
        match (codec, self) {
            (Codec::Hevc, _) => true,
            (Codec::Av1, Self::Nvenc | Self::Qsv | Self::Amf | Self::Software) => true,
            (Codec::Av1, Self::VideoToolbox) => false,
        }
    }

    pub fn codec_name(self, codec: Codec) -> &'static str {
        match (codec, self) {
            (Codec::Hevc, Self::Nvenc) => "hevc_nvenc",
            (Codec::Hevc, Self::VideoToolbox) => "hevc_videotoolbox",
            (Codec::Hevc, Self::Qsv) => "hevc_qsv",
            (Codec::Hevc, Self::Amf) => "hevc_amf",
            (Codec::Hevc, Self::Software) => "libx265",
            (Codec::Av1, Self::Nvenc) => "av1_nvenc",
            (Codec::Av1, Self::Qsv) => "av1_qsv",
            (Codec::Av1, Self::Amf) => "av1_amf",
            (Codec::Av1, Self::Software) => "libsvtav1",
            (Codec::Av1, Self::VideoToolbox) => unreachable!("unsupported encoder/codec pair"),
        }
    }

    pub fn input_args(self) -> &'static [&'static str] {
        match self {
            Self::Nvenc => &["-hwaccel", "cuda"],
            _ => &[],
        }
    }

    /// Maps shrivel's 0-51 compression scale onto this backend's native range.
    pub fn video_args(self, codec: Codec, cq: u8) -> Vec<String> {
        let q = native_quality(self, codec, cq).to_string();
        let args: Vec<String> = match (codec, self) {
            (Codec::Hevc, Self::Nvenc) => vec![
                "-c:v".into(),
                "hevc_nvenc".into(),
                "-preset".into(),
                "p7".into(),
                "-tune".into(),
                "hq".into(),
                "-rc".into(),
                "vbr".into(),
                "-cq".into(),
                q.clone(),
                "-b:v".into(),
                "0".into(),
                "-spatial-aq".into(),
                "1".into(),
                "-temporal-aq".into(),
                "1".into(),
                "-aq-strength".into(),
                "4".into(),
                "-bf".into(),
                "3".into(),
                "-b_ref_mode".into(),
                "middle".into(),
                "-rc-lookahead".into(),
                "32".into(),
            ],
            (Codec::Hevc, Self::Qsv) => vec![
                "-c:v".into(),
                "hevc_qsv".into(),
                "-preset".into(),
                "veryslow".into(),
                "-global_quality".into(),
                q.clone(),
                "-look_ahead".into(),
                "1".into(),
            ],
            (Codec::Hevc, Self::Amf) => vec![
                "-c:v".into(),
                "hevc_amf".into(),
                "-quality".into(),
                "quality".into(),
                "-rc".into(),
                "cqp".into(),
                "-qp_i".into(),
                q.clone(),
                "-qp_p".into(),
                q.clone(),
            ],
            (Codec::Hevc, Self::VideoToolbox) => vec![
                "-c:v".into(),
                "hevc_videotoolbox".into(),
                "-q:v".into(),
                videotoolbox_quality(cq).to_string(),
            ],
            (Codec::Hevc, Self::Software) => vec![
                "-c:v".into(),
                "libx265".into(),
                "-preset".into(),
                "slow".into(),
                "-crf".into(),
                q.clone(),
            ],
            (Codec::Av1, Self::Nvenc) => vec![
                "-c:v".into(),
                "av1_nvenc".into(),
                "-preset".into(),
                "p7".into(),
                "-tune".into(),
                "hq".into(),
                "-rc".into(),
                "vbr".into(),
                "-cq".into(),
                q.clone(),
                "-b:v".into(),
                "0".into(),
            ],
            (Codec::Av1, Self::Qsv) => vec![
                "-c:v".into(),
                "av1_qsv".into(),
                "-preset".into(),
                "veryslow".into(),
                "-global_quality".into(),
                q.clone(),
            ],
            (Codec::Av1, Self::Amf) => vec![
                "-c:v".into(),
                "av1_amf".into(),
                "-quality".into(),
                "quality".into(),
                "-rc".into(),
                "cqp".into(),
                "-qp_i".into(),
                q.clone(),
                "-qp_p".into(),
                q.clone(),
            ],
            (Codec::Av1, Self::Software) => vec![
                "-c:v".into(),
                "libsvtav1".into(),
                "-preset".into(),
                "6".into(),
                "-crf".into(),
                q.clone(),
            ],
            (Codec::Av1, Self::VideoToolbox) => unreachable!("unsupported encoder/codec pair"),
        };
        args
    }

    pub fn is_usable(self, codec: Codec) -> bool {
        let mut cmd = Command::new("ffmpeg");
        cmd.args([
            "-hide_banner",
            "-loglevel",
            "error",
            "-nostdin",
            "-f",
            "lavfi",
            "-i",
            "color=c=black:s=640x360:r=1:d=1",
            "-frames:v",
            "1",
        ]);
        cmd.args(self.video_args(codec, 30))
            .args(["-f", "null", "-"]);
        cmd.stdin(Stdio::null())
            .stdout(Stdio::null())
            .stderr(Stdio::null());
        cmd.status().map(|s| s.success()).unwrap_or(false)
    }

    pub fn detect(codec: Codec) -> Self {
        let candidates = match codec {
            Codec::Hevc => &Self::HEVC_HARDWARE[..],
            Codec::Av1 => &Self::AV1_HARDWARE[..],
        };
        candidates
            .iter()
            .copied()
            .find(|e| e.is_usable(codec))
            .unwrap_or(Self::Software)
    }
}

fn videotoolbox_quality(cq: u8) -> u8 {
    100 - scale_linear(cq, 0, 99)
}

/// Native values are different measurements; this maps their usable ranges,
/// not perceptually equivalent quality. Avoid NVENC's special `0 = auto`.
fn native_quality(encoder: Encoder, codec: Codec, cq: u8) -> u8 {
    match (codec, encoder) {
        (Codec::Hevc, Encoder::Software) => cq, // x265 CRF: 0..51
        (Codec::Av1, Encoder::Software) => scale_linear(cq, 0, 63), // SVT-AV1 CRF
        (Codec::Hevc, Encoder::Nvenc) => scale_linear(cq, 1, 51),
        (Codec::Av1, Encoder::Nvenc) => scale_linear(cq, 1, 63),
        (_, Encoder::Qsv) => scale_linear(cq, 1, 51), // QSV ICQ: 1 is best
        (Codec::Hevc, Encoder::Amf) => cq,            // HEVC QP: 0..51
        (Codec::Av1, Encoder::Amf) => scale_linear(cq, 0, 255), // AV1 QP
        (Codec::Hevc, Encoder::VideoToolbox) => videotoolbox_quality(cq),
        (Codec::Av1, Encoder::VideoToolbox) => unreachable!("unsupported encoder/codec pair"),
    }
}

fn scale_linear(value: u8, min: u8, max: u8) -> u8 {
    let numerator = value.min(51) as u16 * (max - min) as u16;
    (min + ((numerator + 25) / 51) as u8).min(max)
}