shrivel 0.1.3

//! Source inspection through `ffprobe`.

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

use anyhow::{Context, Result, bail};
use serde::Deserialize;

/// The properties of a video stream that drive conversion decisions.
#[derive(Debug, Clone, PartialEq)]
pub struct VideoInfo {
    /// Codec name as reported by FFmpeg (for example `h264` or `hevc`).
    pub codec: String,
    /// Bitrate in kbit/s, or 0 when it cannot be determined.
    pub bitrate_kbps: u32,
    /// Frame rate in frames per second, or 0.0 when unknown.
    pub fps: f64,
    /// Duration in seconds, or 0.0 when unknown.
    pub duration: f64,
}

#[derive(Deserialize)]
struct RawProbe {
    #[serde(default)]
    streams: Vec<RawStream>,
    format: Option<RawFormat>,
}

#[derive(Deserialize)]
struct RawStream {
    codec_name: Option<String>,
    bit_rate: Option<String>,
    r_frame_rate: Option<String>,
    duration: Option<String>,
}

#[derive(Deserialize)]
struct RawFormat {
    bit_rate: Option<String>,
    duration: Option<String>,
}

/// Runs `ffprobe` on `path` and extracts the first video stream's properties.
pub fn probe(path: &Path) -> Result<VideoInfo> {
    let output = Command::new("ffprobe")
        .args([
            "-v",
            "error",
            "-select_streams",
            "v:0",
            "-show_entries",
            "stream=codec_name,bit_rate,r_frame_rate,duration:format=bit_rate,duration",
            "-of",
            "json",
        ])
        .arg(path)
        .stdin(Stdio::null())
        .output()
        .context("failed to run ffprobe")?;

    if !output.status.success() {
        bail!(
            "ffprobe failed: {}",
            String::from_utf8_lossy(&output.stderr).trim()
        );
    }
    parse(&output.stdout)
}

/// Parses `ffprobe -of json` output.
fn parse(json: &[u8]) -> Result<VideoInfo> {
    let raw: RawProbe = serde_json::from_slice(json).context("unreadable ffprobe output")?;
    let stream = raw
        .streams
        .into_iter()
        .next()
        .context("no video stream found")?;

    // Some containers (e.g. Matroska) do not report a per-stream bitrate, so
    // fall back to the container-level value.
    let bits_per_second = stream
        .bit_rate
        .as_deref()
        .and_then(parse_u64)
        .or_else(|| {
            raw.format
                .as_ref()
                .and_then(|f| f.bit_rate.as_deref())
                .and_then(parse_u64)
        })
        .unwrap_or(0);

    Ok(VideoInfo {
        codec: stream.codec_name.unwrap_or_default(),
        bitrate_kbps: (bits_per_second as f64 / 1000.0).round() as u32,
        fps: stream
            .r_frame_rate
            .as_deref()
            .and_then(parse_rational)
            .unwrap_or(0.0),
        duration: stream
            .duration
            .as_deref()
            .and_then(|s| s.parse().ok())
            .or_else(|| {
                raw.format
                    .as_ref()
                    .and_then(|f| f.duration.as_deref())
                    .and_then(|s| s.parse().ok())
            })
            .unwrap_or(0.0),
    })
}

fn parse_u64(s: &str) -> Option<u64> {
    s.trim().parse().ok()
}

/// Parses a frame rate such as `25/1`, `60000/1001` or `30`, rounded to two
/// decimals. Returns `None` for malformed values or a zero denominator.
fn parse_rational(s: &str) -> Option<f64> {
    let value = match s.trim().split_once('/') {
        Some((num, den)) => {
            let num: f64 = num.parse().ok()?;
            let den: f64 = den.parse().ok()?;
            if den == 0.0 {
                return None;
            }
            num / den
        }
        None => s.trim().parse().ok()?,
    };
    Some((value * 100.0).round() / 100.0)
}

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

    #[test]
    fn rational_frame_rates() {
        assert_eq!(parse_rational("25/1"), Some(25.0));
        assert_eq!(parse_rational("60000/1001"), Some(59.94));
        assert_eq!(parse_rational("30"), Some(30.0));
        assert_eq!(parse_rational("0/0"), None);
        assert_eq!(parse_rational("abc"), None);
    }

    #[test]
    fn parses_stream_bitrate() {
        let json =
            br#"{"streams":[{"codec_name":"h264","bit_rate":"4500000","r_frame_rate":"30/1"}],
                        "format":{"bit_rate":"4700000"}}"#;
        let info = parse(json).unwrap();
        assert_eq!(
            info,
            VideoInfo {
                codec: "h264".into(),
                bitrate_kbps: 4500,
                fps: 30.0,
                duration: 0.0,
            }
        );
    }

    #[test]
    fn falls_back_to_container_bitrate() {
        let json = br#"{"streams":[{"codec_name":"hevc","r_frame_rate":"24/1"}],
                        "format":{"bit_rate":"8000000"}}"#;
        assert_eq!(parse(json).unwrap().bitrate_kbps, 8000);
    }

    #[test]
    fn missing_bitrate_becomes_zero() {
        let json = br#"{"streams":[{"codec_name":"h264","r_frame_rate":"24/1"}]}"#;
        assert_eq!(parse(json).unwrap().bitrate_kbps, 0);
    }

    #[test]
    fn no_video_stream_is_an_error() {
        assert!(parse(br#"{"streams":[]}"#).is_err());
    }
}