shrivel 0.2.1

//! shrivel: a cross-platform FFmpeg orchestrator that batch re-encodes videos
//! to H.265/HEVC.

mod cli;
mod encoder;
mod job;
mod probe;
mod progress;

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

use anyhow::{Context, Result, bail};
use clap::Parser;
use owo_colors::OwoColorize;

use cli::{Cli, Codec};
use encoder::Encoder;
use job::Plan;
use progress::ProgressBars;

/// Maximum number of FFmpeg error lines shown for a failed conversion.
const MAX_ERROR_LINES: usize = 15;

const BYTES_PER_MB: f64 = 1024.0 * 1024.0;

/// Settings that stay constant for the whole batch.
struct Settings {
    codec: Codec,
    encoder: Encoder,
    cq: u8,
    target_fps: Option<f64>,
    denoise: bool,
    skip_same_codec: bool,
    skip_existing: bool,
    scale: Option<u32>,
    dry_run: bool,
}

/// What happened to a single file.
enum Outcome {
    Converted { input_bytes: u64, output_bytes: u64 },
    Skipped,
    Failed,
}

fn main() -> ExitCode {
    match run() {
        Ok(code) => code,
        Err(err) => {
            eprintln!("{} {err:#}", "[ERROR]".red());
            ExitCode::FAILURE
        }
    }
}

fn run() -> Result<ExitCode> {
    let cli = Cli::parse();

    require_tool("ffmpeg")?;
    require_tool("ffprobe")?;

    let input_is_file = cli.input.is_file();
    if !input_is_file && !cli.input.is_dir() {
        bail!(
            "input path is not a file or directory: {}",
            cli.input.display()
        );
    }

    let extensions: Vec<String> = cli
        .ext
        .iter()
        .map(|e| e.trim().trim_start_matches('.').to_ascii_lowercase())
        .filter(|e| !e.is_empty())
        .collect();
    if extensions.is_empty() {
        bail!("no valid file extensions given to --ext");
    }

    let (files, input_root) = if input_is_file {
        let matches = cli
            .input
            .extension()
            .and_then(|e| e.to_str())
            .is_some_and(|ext| extensions.iter().any(|x| x.eq_ignore_ascii_case(ext)));
        (
            if matches {
                vec![cli.input.clone()]
            } else {
                vec![]
            },
            cli.input.parent().unwrap_or(Path::new(".")).to_path_buf(),
        )
    } else {
        (collect_inputs(&cli.input, &extensions)?, cli.input.clone())
    };
    if files.is_empty() {
        println!(
            "{}",
            format!(
                "No files with extension(s) {} found.",
                extensions.join(", ")
            )
            .yellow()
        );
        return Ok(ExitCode::SUCCESS);
    }

    if !cli.dry_run {
        fs::create_dir_all(&cli.output)
            .with_context(|| format!("cannot create output directory {}", cli.output.display()))?;
    }
    if input_is_file {
        ensure_distinct_file(
            &files[0],
            &cli.output.join(files[0].file_name().unwrap_or_default()),
        )?;
    } else {
        ensure_distinct_dirs(&cli.input, &cli.output)?;
    }

    let encoder = resolve_encoder(&cli)?;
    let settings = Settings {
        encoder,
        cq: cli.cq,
        target_fps: cli.fps,
        denoise: cli.denoise,
        codec: cli.codec,
        skip_same_codec: cli.skip_same_codec,
        skip_existing: cli.skip_existing,
        scale: cli.scale,
        dry_run: cli.dry_run,
    };

    println!(
        "{}",
        format!(
            "Found {} file(s). Codec={}, Encoder={}, CQ={}, TargetFPS={}, SkipSameCodec={}, Denoise={}, Scale={}",
            files.len(),
            settings.codec,
            encoder.codec_name(settings.codec),
            settings.cq,
            settings
                .target_fps
                .map_or_else(|| "source".to_string(), |f| f.to_string()),
            settings.skip_same_codec,
            settings.denoise,
            settings
                .scale
                .map_or_else(|| "source".to_string(), |w| w.to_string()),
        )
        .cyan()
    );

    let (mut done, mut failed, mut skipped) = (0u32, 0u32, 0u32);
    let (mut total_in, mut total_out) = (0u64, 0u64);

    let progress = ProgressBars::new(files.len() as u64);

    for file in &files {
        let file_name = file.file_name().unwrap_or_default().to_string_lossy();
        progress.set_message(format!("Converting {}", file_name));
        match convert_one(
            &settings,
            file,
            &input_root,
            &cli.output,
            &progress,
            &file_name,
        ) {
            Outcome::Converted {
                input_bytes,
                output_bytes,
            } => {
                done += 1;
                total_in += input_bytes;
                total_out += output_bytes;
            }
            Outcome::Skipped => skipped += 1,
            Outcome::Failed => failed += 1,
        }
        progress.advance();
    }
    progress.clear()?;

    let done_label = if settings.dry_run {
        "Planned"
    } else {
        "Succeeded"
    };
    println!("\n{}", "=== Summary ===".yellow());
    println!(
        "{}",
        format!("{done_label}: {done} - Failed: {failed} - Skipped: {skipped}").yellow()
    );
    if total_in > 0 {
        println!(
            "{}",
            format!(
                "Total: {:.1} MB -> {:.1} MB ({:.1}%)",
                total_in as f64 / BYTES_PER_MB,
                total_out as f64 / BYTES_PER_MB,
                total_out as f64 / total_in as f64 * 100.0
            )
            .yellow()
        );
    }

    Ok(if failed > 0 {
        ExitCode::FAILURE
    } else {
        ExitCode::SUCCESS
    })
}

/// Fails with a friendly message when an external tool cannot be started.
fn require_tool(name: &str) -> Result<()> {
    Command::new(name)
        .arg("-version")
        .stdin(Stdio::null())
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .status()
        .map(|_| ())
        .with_context(|| format!("`{name}` was not found in PATH"))
}

/// Refuses to run when input and output are the same directory, which would
/// make FFmpeg read and overwrite the same file.
fn ensure_distinct_dirs(input: &Path, output: &Path) -> Result<()> {
    if let (Ok(a), Ok(b)) = (fs::canonicalize(input), fs::canonicalize(output))
        && a == b
    {
        bail!("input and output directories must be different");
    }
    Ok(())
}

fn ensure_distinct_file(input: &Path, output: &Path) -> Result<()> {
    if let (Ok(a), Ok(b)) = (fs::canonicalize(input), fs::canonicalize(output))
        && a == b
    {
        bail!("input and output files must be different");
    }
    Ok(())
}

/// Lists files recursively whose extension matches one of `extensions`
/// (case-insensitive), sorted by path.
fn collect_inputs(dir: &Path, extensions: &[String]) -> Result<Vec<PathBuf>> {
    let mut files = Vec::new();
    collect_inputs_recursive(dir, extensions, &mut files)?;
    files.sort();
    Ok(files)
}

fn collect_inputs_recursive(
    dir: &Path,
    extensions: &[String],
    files: &mut Vec<PathBuf>,
) -> Result<()> {
    let entries =
        fs::read_dir(dir).with_context(|| format!("cannot read directory {}", dir.display()))?;
    for entry in entries {
        let path = entry?.path();
        let file_type = fs::symlink_metadata(&path)
            .with_context(|| format!("cannot inspect {}", path.display()))?
            .file_type();
        if file_type.is_dir() {
            collect_inputs_recursive(&path, extensions, files)?;
        } else if file_type.is_file()
            && path
                .extension()
                .and_then(|e| e.to_str())
                .is_some_and(|ext| extensions.iter().any(|x| x.eq_ignore_ascii_case(ext)))
        {
            files.push(path);
        }
    }
    Ok(())
}

fn output_path(input: &Path, input_root: &Path, output_root: &Path) -> Option<PathBuf> {
    input
        .strip_prefix(input_root)
        .ok()
        .map(|relative_path| output_root.join(relative_path))
}

/// Resolves the encoder from the CLI choice, probing the machine as needed.
fn resolve_encoder(cli: &Cli) -> Result<Encoder> {
    if cli.encoder == cli::EncoderChoice::Videotoolbox && cli.codec == Codec::Av1 {
        bail!(
            "encoder `videotoolbox` does not support AV1; choose auto, nvenc, qsv, amf or software"
        );
    }
    match Encoder::from_choice(cli.encoder, cli.codec) {
        Some(encoder) => {
            // Dry runs never execute FFmpeg for encoding, so skip the probe.
            if !cli.dry_run && !encoder.is_usable(cli.codec) {
                bail!(
                    "encoder `{}` is not usable on this machine \
                     (check your FFmpeg build and drivers, or use --encoder auto)",
                    encoder.codec_name(cli.codec)
                );
            }
            Ok(encoder)
        }
        None => {
            let encoder = Encoder::detect(cli.codec);
            println!(
                "{}",
                format!("Auto-detected encoder: {}", encoder.codec_name(cli.codec)).bright_black()
            );
            Ok(encoder)
        }
    }
}

/// Converts one file and reports the result on the terminal.
fn convert_one(
    settings: &Settings,
    input: &Path,
    input_root: &Path,
    output_dir: &Path,
    progress: &ProgressBars,
    file_name: &std::borrow::Cow<'_, str>,
) -> Outcome {
    let output = match output_path(input, input_root, output_dir) {
        Some(path) => path,
        None => {
            progress.println(format!(
                "     {} cannot determine path relative to input",
                "[ERR]".red()
            ));
            return Outcome::Failed;
        }
    };

    if !settings.dry_run
        && let Some(parent) = output.parent()
        && let Err(err) = fs::create_dir_all(parent)
    {
        progress.println(format!(
            "     {} cannot create output directory {}: {err}",
            "[ERR]".red(),
            parent.display()
        ));
        return Outcome::Failed;
    }

    progress.println(format!("\n{} {}", "[>>]".cyan(), file_name.cyan()));

    let info = match probe::probe(input) {
        Ok(info) => info,
        Err(err) => {
            progress.println(format!("     {} {err:#}", "[ERR]".red()));
            return Outcome::Failed;
        }
    };
    progress.println(
        format!(
            "     Codec: {} | Bitrate: {} kbps | FPS: {} | Duration: {:.1}s",
            info.codec, info.bitrate_kbps, info.fps, info.duration
        )
        .bright_black()
        .to_string(),
    );

    let target_codec_name = match settings.codec {
        Codec::Hevc => "hevc",
        Codec::Av1 => "av1",
    };
    if settings.skip_same_codec && info.codec == target_codec_name {
        progress.println(format!(
            "     {}",
            format!("[SKIP] Already {target_codec_name}.").yellow()
        ));
        return Outcome::Skipped;
    }
    if settings.skip_existing && output.exists() {
        progress.println(format!("     {}", "[SKIP] Output already exists.".yellow()));
        return Outcome::Skipped;
    }

    let cq = settings.cq;
    progress.println(
        format!("     CQ: {cq} (normalized)")
            .bright_black()
            .to_string(),
    );

    let mut filters = Vec::new();
    let mut force_cfr = false;
    if settings.denoise {
        filters.push("hqdn3d=3:3:8:8".to_string());
    }
    if let Some(scale) = settings.scale {
        // scale=width:-2 ensures the height is proportional and divisible by 2.
        // We explicitly use bicubic interpolation to ensure consistent behavior.
        filters.push(format!("scale={scale}:-2:flags=bicubic"));
    }
    if let Some(target) = settings.target_fps {
        if job::fps_change_needed(info.fps, target) {
            filters.push(format!("fps={target}"));
            // Constant frame rate avoids leftover variable-frame-rate drift.
            force_cfr = true;
            progress.println(
                format!("     FPS: {} -> {target}", info.fps)
                    .bright_black()
                    .to_string(),
            );
        } else {
            progress.println(
                format!("     FPS: already {}, no change", info.fps)
                    .bright_black()
                    .to_string(),
            );
        }
    }

    let plan = Plan {
        codec: settings.codec,
        encoder: settings.encoder,
        input,
        output: &output,
        cq,
        filters,
        force_cfr,
    };
    let args = plan.ffmpeg_args();

    if settings.dry_run {
        let rendered: Vec<String> = args.iter().map(|a| quote(&a.to_string_lossy())).collect();
        println!(
            "{}",
            format!("     {} ffmpeg {}", "[DRY]".magenta(), rendered.join(" "))
        );
        return Outcome::Converted {
            input_bytes: 0,
            output_bytes: 0,
        };
    }

    let mut child = match Command::new("ffmpeg")
        .args(&args)
        .stdin(Stdio::null())
        .stderr(Stdio::piped())
        .spawn()
    {
        Ok(child) => child,
        Err(err) => {
            progress.println(format!(
                "     {} could not start ffmpeg: {err}",
                "[ERR]".red()
            ));
            return Outcome::Failed;
        }
    };

    let mut stderr = match child.stderr.take() {
        Some(s) => s,
        None => {
            progress.println(format!("     {} failed to take stderr", "[ERR]".red()));
            return Outcome::Failed;
        }
    };
    // Keep a reusable buffer with writable space for each stderr read.
    let mut buffer = [0u8; 8192];
    let mut total_err_output = String::new();
    let mut pending_progress = String::new();

    let file_progress = progress.file(info.duration);

    use std::io::Read;
    loop {
        let bytes_read = stderr.read(&mut buffer).unwrap_or(0);
        if bytes_read == 0 {
            break;
        }
        let chunk = String::from_utf8_lossy(&buffer[..bytes_read]);
        total_err_output.push_str(&chunk);
        pending_progress.push_str(&chunk);

        // `-progress pipe:2` emits machine-readable `out_time=HH:MM:SS...`
        // lines, independently of FFmpeg's log level.
        while let Some(newline) = pending_progress.find('\n') {
            let line = pending_progress[..newline].trim().to_owned();
            pending_progress.drain(..=newline);
            if let Some(time) = line.strip_prefix("out_time=") {
                if let Some(seconds) = parse_ffmpeg_time(time.trim()) {
                    if let Some(ref p) = file_progress {
                        p.set_seconds(seconds);
                    }
                }
            }
        }
    }

    if let Some(p) = file_progress {
        p.finish();
    }

    let status = match child.wait() {
        Ok(s) => s,
        Err(err) => {
            progress.println(format!(
                "     {} failed to wait for ffmpeg: {err}",
                "[ERR]".red()
            ));
            return Outcome::Failed;
        }
    };

    if status.success() {
        let input_bytes = fs::metadata(input).map(|m| m.len()).unwrap_or(0);
        let output_bytes = fs::metadata(&output).map(|m| m.len()).unwrap_or(0);
        let ratio = if input_bytes > 0 {
            output_bytes as f64 / input_bytes as f64 * 100.0
        } else {
            0.0
        };
        let line = format!(
            "     [OK] {:.1} MB -> {:.1} MB ({:.1}%)",
            input_bytes as f64 / BYTES_PER_MB,
            output_bytes as f64 / BYTES_PER_MB,
            ratio
        );
        if ratio <= 100.0 {
            progress.println(line.green().to_string());
        } else {
            progress.println(line.yellow().to_string());
        }
        Outcome::Converted {
            input_bytes,
            output_bytes,
        }
    } else {
        progress.println(format!("     {}", "[ERR] Conversion failed".red()));
        let lines: Vec<_> = total_err_output.lines().collect();
        for line in lines.iter().rev().take(MAX_ERROR_LINES).rev() {
            progress.println(format!("     {}", line.bright_black()));
        }
        let _ = fs::remove_file(&output);
        Outcome::Failed
    }
}

fn parse_ffmpeg_time(s: &str) -> Option<f64> {
    let parts: Vec<&str> = s.split(':').collect();
    if parts.len() != 3 {
        return None;
    }
    let h: f64 = parts[0].parse().ok()?;
    let m: f64 = parts[1].parse().ok()?;
    let s: f64 = parts[2].parse().ok()?;
    Some(h * 3600.0 + m * 60.0 + s)
}

/// Quotes an argument for display only, so printed dry-run commands can be
/// pasted into a shell. The command itself is never run through a shell.
fn quote(arg: &str) -> String {
    let needs_quotes = arg.is_empty()
        || arg
            .chars()
            .any(|c| c.is_whitespace() || "\"'$&|;<>()*?[]{}!\\`".contains(c));
    if needs_quotes {
        format!("\"{}\"", arg.replace('"', "\\\""))
    } else {
        arg.to_string()
    }
}

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

    #[test]
    fn quote_leaves_plain_arguments_alone() {
        assert_eq!(quote("-hide_banner"), "-hide_banner");
        assert_eq!(quote("0:a?"), "\"0:a?\"");
    }

    #[test]
    fn quote_wraps_arguments_with_spaces() {
        assert_eq!(quote("my video.mp4"), "\"my video.mp4\"");
        assert_eq!(quote(""), "\"\"");
    }

    #[test]
    fn collect_inputs_filters_by_extension_case_insensitively() {
        let dir = std::env::temp_dir().join(format!("shrivel-test-{}", std::process::id()));
        fs::create_dir_all(dir.join("nested/deeper")).unwrap();
        for name in ["b.MP4", "a.mp4", "c.mkv", "d.txt"] {
            fs::write(dir.join(name), b"").unwrap();
        }
        fs::write(dir.join("nested/video.mp4"), b"").unwrap();
        fs::write(dir.join("nested/deeper/other.MP4"), b"").unwrap();
        let found = collect_inputs(&dir, &["mp4".to_string()]).unwrap();
        let relative_paths: Vec<Vec<String>> = found
            .iter()
            .map(|p| {
                p.strip_prefix(&dir)
                    .unwrap()
                    .components()
                    .map(|component| component.as_os_str().to_string_lossy().into_owned())
                    .collect()
            })
            .collect();
        assert_eq!(
            relative_paths,
            [
                vec!["a.mp4"],
                vec!["b.MP4"],
                vec!["nested", "deeper", "other.MP4"],
                vec!["nested", "video.mp4"]
            ]
        );
        fs::remove_dir_all(&dir).unwrap();
    }

    #[test]
    fn output_path_preserves_input_subdirectories() {
        let input_root = Path::new("input");
        let output_root = Path::new("output");
        let input = input_root.join("nested/deeper/video.mp4");

        assert_eq!(
            output_path(&input, input_root, output_root),
            Some(output_root.join("nested/deeper/video.mp4"))
        );
    }
}