shrivel 0.4.0

//! 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,
    pub(crate) duration_origin: DurationOrigin,
    /// Width of the video stream in pixels.
    pub width: u32,
    /// Height of the video stream in pixels.
    pub height: u32,
}

/// The ffprobe field used for the reported duration.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum DurationOrigin {
    Stream,
    Format,
    Unknown,
}

#[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>,
    width: Option<u32>,
    height: Option<u32>,
}

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

#[derive(Deserialize)]
struct StreamPresenceProbe {
    #[serde(default)]
    streams: Vec<StreamPresence>,
}

#[derive(Deserialize)]
struct StreamPresence {
    codec_type: Option<String>,
    disposition: Option<StreamDisposition>,
}

#[derive(Deserialize)]
struct StreamDisposition {
    attached_pic: Option<u8>,
}

/// 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,width,height: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)
}

/// Returns whether ffprobe finds actual video, excluding attached cover art.
pub fn has_video_stream(path: &Path) -> Result<bool> {
    let output = Command::new("ffprobe")
        .args([
            "-v",
            "error",
            "-show_entries",
            "stream=codec_type:stream_disposition=attached_pic",
            "-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_has_video_stream(&output.stdout)
}

fn parse_has_video_stream(json: &[u8]) -> Result<bool> {
    let probe: StreamPresenceProbe =
        serde_json::from_slice(json).context("unreadable ffprobe stream output")?;
    Ok(probe.streams.iter().any(|stream| {
        stream.codec_type.as_deref() == Some("video")
            && stream
                .disposition
                .as_ref()
                .and_then(|d| d.attached_pic)
                .unwrap_or(0)
                == 0
    }))
}

/// 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);

    let (duration, duration_origin) = stream
        .duration
        .as_deref()
        .and_then(parse_duration)
        .map(|duration| (duration, DurationOrigin::Stream))
        .or_else(|| {
            raw.format
                .as_ref()
                .and_then(|f| f.duration.as_deref())
                .and_then(parse_duration)
                .map(|duration| (duration, DurationOrigin::Format))
        })
        .unwrap_or((0.0, DurationOrigin::Unknown));

    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,
        duration_origin,
        width: stream.width.unwrap_or(0),
        height: stream.height.unwrap_or(0),
    })
}

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

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

fn parse_duration(s: &str) -> Option<f64> {
    parse_f64(s).filter(|duration| duration.is_finite() && *duration > 0.0)
}

/// 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 attached_cover_art_does_not_count_as_video() {
        let json = br#"{"streams":[{"codec_type":"audio"},{"codec_type":"video","disposition":{"attached_pic":1}}]}"#;
        assert!(!parse_has_video_stream(json).unwrap());
    }

    #[test]
    fn real_video_counts_even_when_cover_art_is_present() {
        let json = br#"{"streams":[{"codec_type":"audio"},{"codec_type":"video","disposition":{"attached_pic":1}},{"codec_type":"video","disposition":{"attached_pic":0}}]}"#;
        assert!(parse_has_video_stream(json).unwrap());
    }

    #[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,
                duration_origin: DurationOrigin::Unknown,
                width: 0,
                height: 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 stream_duration_takes_precedence_over_format_duration() {
        let json = br#"{"streams":[{"duration":"12.5"}],"format":{"duration":"13.0"}}"#;
        let info = parse(json).unwrap();
        assert_eq!(info.duration, 12.5);
        assert_eq!(info.duration_origin, DurationOrigin::Stream);
    }

    #[test]
    fn falls_back_to_valid_format_duration() {
        let json = br#"{"streams":[{}],"format":{"duration":"13.0"}}"#;
        let info = parse(json).unwrap();
        assert_eq!(info.duration, 13.0);
        assert_eq!(info.duration_origin, DurationOrigin::Format);
    }

    #[test]
    fn invalid_durations_become_unknown() {
        let json = br#"{"streams":[{"duration":"NaN"}],"format":{"duration":"0"}}"#;
        let info = parse(json).unwrap();
        assert_eq!(info.duration, 0.0);
        assert_eq!(info.duration_origin, DurationOrigin::Unknown);
    }

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