noai 0.1.0

A small CLI that re-encodes images through JPEG and PNG with noai-prefixed output names.
use std::env;
use std::error::Error;
use std::ffi::{OsStr, OsString};
use std::fmt;
use std::fs;
use std::path::{Path, PathBuf};

use image::codecs::jpeg::JpegEncoder;
use image::{DynamicImage, ImageFormat};

type AppResult<T> = Result<T, Box<dyn Error>>;

const OUTPUT_PREFIX: &str = "noai_";
const USAGE: &str = "\
Usage:
  noai image <file-or-directory>
  noai image -r <directory>
";

#[derive(Debug)]
struct CliError(String);

impl fmt::Display for CliError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.0)
    }
}

impl Error for CliError {}

#[derive(Debug, PartialEq, Eq)]
struct ImageCommand {
    recursive: bool,
    path: PathBuf,
}

fn main() {
    if let Err(error) = run(env::args_os().skip(1)) {
        eprintln!("{error}");
        std::process::exit(1);
    }
}

fn run<I>(args: I) -> AppResult<()>
where
    I: IntoIterator<Item = OsString>,
{
    let command = parse_args(args)?;

    if command.recursive && !command.path.is_dir() {
        return Err(cli_error(format!(
            "-r requires a directory: {}\n{USAGE}",
            command.path.display()
        )));
    }

    let image_files = collect_image_files(&command.path, command.recursive)?;

    if image_files.is_empty() {
        return Err(cli_error(format!(
            "No supported image files found in {}\n{USAGE}",
            command.path.display()
        )));
    }

    let mut failures = Vec::new();

    for source in image_files {
        match transcode_file(&source) {
            Ok(output) => println!("{} -> {}", source.display(), output.display()),
            Err(error) => failures.push(format!("{}: {error}", source.display())),
        }
    }

    if failures.is_empty() {
        Ok(())
    } else {
        Err(cli_error(format!(
            "Failed to convert {} file(s):\n{}",
            failures.len(),
            failures.join("\n")
        )))
    }
}

fn parse_args<I>(args: I) -> AppResult<ImageCommand>
where
    I: IntoIterator<Item = OsString>,
{
    let mut args = args.into_iter();

    match args.next() {
        Some(command) if command == OsStr::new("image") => {}
        Some(command) => {
            return Err(cli_error(format!(
                "Unknown command: {}\n{USAGE}",
                command.to_string_lossy()
            )));
        }
        None => return Err(cli_error(format!("Missing command\n{USAGE}"))),
    }

    let first_arg = args
        .next()
        .ok_or_else(|| cli_error(format!("Missing path\n{USAGE}")))?;

    let (recursive, path_arg) = if first_arg == OsStr::new("-r") {
        let path = args
            .next()
            .ok_or_else(|| cli_error(format!("Missing directory after -r\n{USAGE}")))?;
        (true, path)
    } else {
        (false, first_arg)
    };

    if let Some(extra) = args.next() {
        return Err(cli_error(format!(
            "Unexpected argument: {}\n{USAGE}",
            extra.to_string_lossy()
        )));
    }

    let path = PathBuf::from(path_arg);

    Ok(ImageCommand { recursive, path })
}

fn collect_image_files(path: &Path, recursive: bool) -> AppResult<Vec<PathBuf>> {
    if path.is_file() {
        return if is_supported_image_path(path) {
            Ok(vec![path.to_path_buf()])
        } else {
            Err(cli_error(format!(
                "Unsupported image file extension: {}\n{USAGE}",
                path.display()
            )))
        };
    }

    if !path.is_dir() {
        return Err(cli_error(format!(
            "Path does not exist: {}\n{USAGE}",
            path.display()
        )));
    }

    let mut files = Vec::new();
    collect_from_directory(path, recursive, &mut files)?;
    files.sort();
    Ok(files)
}

fn collect_from_directory(
    directory: &Path,
    recursive: bool,
    files: &mut Vec<PathBuf>,
) -> AppResult<()> {
    for entry in fs::read_dir(directory)? {
        let entry = entry?;
        let path = entry.path();

        if path.is_dir() {
            if recursive {
                collect_from_directory(&path, recursive, files)?;
            }
            continue;
        }

        if is_generated_output_path(&path) {
            continue;
        }

        if is_supported_image_path(&path) {
            files.push(path);
        }
    }

    Ok(())
}

fn transcode_file(source: &Path) -> AppResult<PathBuf> {
    let original_image = image::open(source)?;
    let jpeg_bytes = encode_as_high_quality_jpeg(&original_image)?;
    let jpeg_image = image::load_from_memory_with_format(&jpeg_bytes, ImageFormat::Jpeg)?;
    let output = next_available_output_path(source)?;

    jpeg_image.save_with_format(&output, ImageFormat::Png)?;

    Ok(output)
}

fn encode_as_high_quality_jpeg(image: &DynamicImage) -> AppResult<Vec<u8>> {
    let rgb_image = DynamicImage::ImageRgb8(image.to_rgb8());
    let mut bytes = Vec::new();
    let mut encoder = JpegEncoder::new_with_quality(&mut bytes, 100);
    encoder.encode_image(&rgb_image)?;
    Ok(bytes)
}

fn next_available_output_path(source: &Path) -> AppResult<PathBuf> {
    let directory = source.parent().unwrap_or_else(|| Path::new("."));
    let file_stem = source
        .file_stem()
        .and_then(OsStr::to_str)
        .ok_or_else(|| cli_error(format!("Invalid source file name: {}", source.display())))?;

    let base_name = format!("{OUTPUT_PREFIX}{file_stem}");
    let mut candidate = directory.join(format!("{base_name}.png"));

    for index in 1.. {
        if !candidate.exists() {
            return Ok(candidate);
        }

        candidate = directory.join(format!("{base_name} ({index}).png"));
    }

    unreachable!("the filename suffix loop is unbounded")
}

fn is_supported_image_path(path: &Path) -> bool {
    path.extension()
        .and_then(OsStr::to_str)
        .map(|extension| {
            matches!(
                extension.to_ascii_lowercase().as_str(),
                "jpg"
                    | "jpeg"
                    | "png"
                    | "webp"
                    | "gif"
                    | "bmp"
                    | "tif"
                    | "tiff"
                    | "tga"
                    | "ico"
                    | "qoi"
            )
        })
        .unwrap_or(false)
}

fn is_generated_output_path(path: &Path) -> bool {
    path.file_name()
        .and_then(OsStr::to_str)
        .map(|file_name| file_name.starts_with(OUTPUT_PREFIX))
        .unwrap_or(false)
}

fn cli_error(message: impl Into<String>) -> Box<dyn Error> {
    Box::new(CliError(message.into()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::{SystemTime, UNIX_EPOCH};

    #[test]
    fn parses_single_file_command() {
        let command = parse_args(["image", "photo.png"].into_iter().map(OsString::from)).unwrap();

        assert_eq!(
            command,
            ImageCommand {
                recursive: false,
                path: PathBuf::from("photo.png")
            }
        );
    }

    #[test]
    fn parses_recursive_directory_command() {
        let command =
            parse_args(["image", "-r", "assets"].into_iter().map(OsString::from)).unwrap();

        assert_eq!(
            command,
            ImageCommand {
                recursive: true,
                path: PathBuf::from("assets")
            }
        );
    }

    #[test]
    fn builds_non_overwriting_output_name() -> AppResult<()> {
        let directory = unique_temp_dir("output-name")?;
        let source = directory.join("photo.png");
        let first_output = directory.join("noai_photo.png");
        let second_output = directory.join("noai_photo (1).png");

        fs::write(&source, b"source")?;
        fs::write(&first_output, b"existing")?;
        fs::write(&second_output, b"existing")?;

        let output = next_available_output_path(&source)?;

        fs::remove_dir_all(&directory)?;

        assert_eq!(output.file_name(), Some(OsStr::new("noai_photo (2).png")));
        Ok(())
    }

    #[test]
    fn transcodes_without_replacing_existing_output() -> AppResult<()> {
        let directory = unique_temp_dir("transcode")?;
        let source = directory.join("sample.png");
        let existing_output = directory.join("noai_sample.png");

        let image = image::RgbImage::from_pixel(2, 2, image::Rgb([24, 48, 96]));
        image.save_with_format(&source, ImageFormat::Png)?;
        fs::write(&existing_output, b"keep this file")?;

        let output = transcode_file(&source)?;
        let loaded_output = image::open(&output)?;
        let existing_contents = fs::read(&existing_output)?;

        fs::remove_dir_all(&directory)?;

        assert_eq!(output.file_name(), Some(OsStr::new("noai_sample (1).png")));
        assert_eq!((loaded_output.width(), loaded_output.height()), (2, 2));
        assert_eq!(existing_contents, b"keep this file");
        Ok(())
    }

    fn unique_temp_dir(name: &str) -> AppResult<PathBuf> {
        let timestamp = SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos();
        let directory =
            env::temp_dir().join(format!("noai-{name}-{}-{timestamp}", std::process::id()));
        fs::create_dir(&directory)?;
        Ok(directory)
    }
}