shrivel 0.1.0

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

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

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

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

use cli::Cli;
use encoder::Encoder;
use job::Plan;

/// 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 {
    encoder: Encoder,
    base_cq: u8,
    target_fps: Option<f64>,
    denoise: bool,
    skip_hevc: bool,
    skip_existing: bool,
    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")?;

    if !cli.input.is_dir() {
        bail!("input directory not found: {}", 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");
    }

    if !cli.dry_run {
        fs::create_dir_all(&cli.output)
            .with_context(|| format!("cannot create output directory {}", cli.output.display()))?;
    }
    ensure_distinct_dirs(&cli.input, &cli.output)?;

    let files = collect_inputs(&cli.input, &extensions)?;
    if files.is_empty() {
        println!(
            "{}",
            format!(
                "No files with extension(s) {} found.",
                extensions.join(", ")
            )
            .yellow()
        );
        return Ok(ExitCode::SUCCESS);
    }

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

    println!(
        "{}",
        format!(
            "Found {} file(s). Encoder={}, BaseCQ={}, TargetFPS={}, SkipHEVC={}, Denoise={}",
            files.len(),
            encoder.codec_name(),
            settings.base_cq,
            settings
                .target_fps
                .map_or_else(|| "source".to_string(), |f| f.to_string()),
            settings.skip_hevc,
            settings.denoise,
        )
        .cyan()
    );

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

    for file in &files {
        match convert_one(&settings, file, &cli.output) {
            Outcome::Converted {
                input_bytes,
                output_bytes,
            } => {
                done += 1;
                total_in += input_bytes;
                total_out += output_bytes;
            }
            Outcome::Skipped => skipped += 1,
            Outcome::Failed => failed += 1,
        }
    }

    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(())
}

/// Lists the files in `dir` whose extension matches one of `extensions`
/// (case-insensitive), sorted by path. Subdirectories are not searched.
fn collect_inputs(dir: &Path, extensions: &[String]) -> Result<Vec<PathBuf>> {
    let mut files = Vec::new();
    let entries =
        fs::read_dir(dir).with_context(|| format!("cannot read directory {}", dir.display()))?;
    for entry in entries {
        let path = entry?.path();
        if !path.is_file() {
            continue;
        }
        let matches = path
            .extension()
            .and_then(|e| e.to_str())
            .is_some_and(|ext| extensions.iter().any(|x| x.eq_ignore_ascii_case(ext)));
        if matches {
            files.push(path);
        }
    }
    files.sort();
    Ok(files)
}

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

/// Converts one file and reports the result on the terminal.
fn convert_one(settings: &Settings, input: &Path, output_dir: &Path) -> Outcome {
    let file_name = input.file_name().unwrap_or_default();
    let output = output_dir.join(file_name);

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

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

    if settings.skip_hevc && info.codec == "hevc" {
        println!("     {}", "[SKIP] Already HEVC.".yellow());
        return Outcome::Skipped;
    }
    if settings.skip_existing && output.exists() {
        println!("     {}", "[SKIP] Output already exists.".yellow());
        return Outcome::Skipped;
    }

    let cq = job::effective_cq(settings.base_cq, info.bitrate_kbps);
    println!("{}", format!("     Effective CQ: {cq}").bright_black());

    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(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;
            println!(
                "{}",
                format!("     FPS: {} -> {target}", info.fps).bright_black()
            );
        } else {
            println!(
                "{}",
                format!("     FPS: already {}, no change", info.fps).bright_black()
            );
        }
    }

    let plan = Plan {
        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!("     {} ffmpeg {}", "[DRY]".magenta(), rendered.join(" "));
        return Outcome::Converted {
            input_bytes: 0,
            output_bytes: 0,
        };
    }

    let result = Command::new("ffmpeg")
        .args(&args)
        .stdin(Stdio::null())
        .output();

    match result {
        Ok(out) if out.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 {
                println!("{}", line.green());
            } else {
                println!("{}", line.yellow());
            }
            Outcome::Converted {
                input_bytes,
                output_bytes,
            }
        }
        Ok(out) => {
            println!("     {}", "[ERR] Conversion failed".red());
            for line in String::from_utf8_lossy(&out.stderr)
                .lines()
                .take(MAX_ERROR_LINES)
            {
                println!("     {}", line.bright_black());
            }
            // Do not leave a truncated, unplayable file behind.
            let _ = fs::remove_file(&output);
            Outcome::Failed
        }
        Err(err) => {
            println!("     {} could not start ffmpeg: {err}", "[ERR]".red());
            Outcome::Failed
        }
    }
}

/// 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("sub.mp4")).unwrap();
        for name in ["b.MP4", "a.mp4", "c.mkv", "d.txt"] {
            fs::write(dir.join(name), b"").unwrap();
        }
        let found = collect_inputs(&dir, &["mp4".to_string()]).unwrap();
        let names: Vec<_> = found
            .iter()
            .map(|p| p.file_name().unwrap().to_string_lossy().into_owned())
            .collect();
        assert_eq!(names, ["a.mp4", "b.MP4"]);
        fs::remove_dir_all(&dir).unwrap();
    }
}