wav1c-cli 0.2.0

Command-line AV1 encoder powered by wav1c.
#![forbid(unsafe_code)]

use std::env;
use std::fs::File;
use std::io::Write;
use std::path::Path;
use std::process;

struct CliArgs {
    input: InputMode,
    output_path: String,
    config: wav1c::EncodeConfig,
}

enum InputMode {
    Y4m(String),
    Solid {
        width: u32,
        height: u32,
        y: u8,
        u: u8,
        v: u8,
    },
}

fn parse_bitrate(args: &[String]) -> Option<u64> {
    for i in 0..args.len().saturating_sub(1) {
        if args[i] == "--bitrate" {
            let s = &args[i + 1];
            let (num, mult) = if let Some(n) = s.strip_suffix('k').or_else(|| s.strip_suffix('K')) {
                (n, 1_000u64)
            } else if let Some(n) = s.strip_suffix('m').or_else(|| s.strip_suffix('M')) {
                (n, 1_000_000u64)
            } else {
                (s.as_str(), 1u64)
            };
            return num.parse::<u64>().ok().map(|v| v * mult);
        }
    }
    None
}

fn parse_option<T: std::str::FromStr>(args: &[String], flag: &str) -> Option<T> {
    for i in 0..args.len().saturating_sub(1) {
        if args[i] == flag {
            return args[i + 1].parse().ok();
        }
    }
    None
}

fn strip_options(args: &[String]) -> Vec<String> {
    let option_flags = ["-q", "--keyint", "--bitrate"];
    let mut result = Vec::new();
    let mut i = 0;
    while i < args.len() {
        if option_flags.contains(&args[i].as_str()) {
            i += 2;
        } else {
            result.push(args[i].clone());
            i += 1;
        }
    }
    result
}

fn parse_cli() -> CliArgs {
    let raw_args: Vec<String> = env::args().collect();

    let quality: u8 = parse_option(&raw_args, "-q").unwrap_or(wav1c::DEFAULT_BASE_Q_IDX);
    let keyint: usize = parse_option(&raw_args, "--keyint").unwrap_or(wav1c::DEFAULT_KEYINT);
    let bitrate: Option<u64> = parse_bitrate(&raw_args);

    let args = strip_options(&raw_args);

    if args.len() < 4 {
        print_usage();
        process::exit(1);
    }

    let config = wav1c::EncodeConfig {
        base_q_idx: quality,
        keyint,
        target_bitrate: bitrate,
        ..Default::default()
    };

    if args[1].ends_with(".y4m") {
        if args.len() < 4 || args[2] != "-o" {
            print_usage();
            process::exit(1);
        }
        CliArgs {
            input: InputMode::Y4m(args[1].clone()),
            output_path: args[3].clone(),
            config,
        }
    } else {
        if args.len() < 8 || args[6] != "-o" {
            print_usage();
            process::exit(1);
        }
        CliArgs {
            input: InputMode::Solid {
                width: args[1].parse().unwrap_or_else(|_| {
                    eprintln!("Error: width must be a positive integer");
                    process::exit(1);
                }),
                height: args[2].parse().unwrap_or_else(|_| {
                    eprintln!("Error: height must be a positive integer");
                    process::exit(1);
                }),
                y: args[3].parse().unwrap_or_else(|_| {
                    eprintln!("Error: Y must be 0-255");
                    process::exit(1);
                }),
                u: args[4].parse().unwrap_or_else(|_| {
                    eprintln!("Error: U must be 0-255");
                    process::exit(1);
                }),
                v: args[5].parse().unwrap_or_else(|_| {
                    eprintln!("Error: V must be 0-255");
                    process::exit(1);
                }),
            },
            output_path: args[7].clone(),
            config,
        }
    }
}

fn print_usage() {
    eprintln!("Usage: wav1c <input.y4m> -o <output.ivf> [options]");
    eprintln!("       wav1c <width> <height> <Y> <U> <V> -o <output.ivf> [options]");
    eprintln!();
    eprintln!("Options:");
    eprintln!("  -q <0-255>      Quantizer index (0=best, 255=smallest, default=128)");
    eprintln!("  --keyint <N>    Keyframe interval in frames (default=25)");
    eprintln!("  --bitrate <N>   Target bitrate (e.g., 500k, 2M). Overrides -q");
}

fn main() {
    let cli = parse_cli();

    let frames = match &cli.input {
        InputMode::Y4m(path) => {
            wav1c::y4m::FramePixels::all_from_y4m_file(Path::new(path)).unwrap_or_else(|e| {
                eprintln!("Error reading {}: {}", path, e);
                process::exit(1);
            })
        }
        InputMode::Solid {
            width,
            height,
            y,
            u,
            v,
        } => vec![wav1c::y4m::FramePixels::solid(*width, *height, *y, *u, *v)],
    };

    let width = frames[0].width;
    let height = frames[0].height;
    let encoder_config = wav1c::EncoderConfig::from(&cli.config);
    let mut encoder = wav1c::Encoder::new(width, height, encoder_config).unwrap_or_else(|e| {
        eprintln!("Error creating encoder: {:?}", e);
        process::exit(1);
    });

    let mut output = Vec::new();
    wav1c::ivf::write_ivf_header(&mut output, width as u16, height as u16, frames.len() as u32)
        .unwrap();

    for frame in &frames {
        encoder.send_frame(frame).unwrap_or_else(|e| {
            eprintln!("Error encoding frame: {:?}", e);
            process::exit(1);
        });

        if let Some(packet) = encoder.receive_packet() {
            let frame_type_str = match packet.frame_type {
                wav1c::FrameType::Key => "KEY",
                wav1c::FrameType::Inter => "INTER",
            };
            eprintln!(
                "frame {:>4}  {:>5}  {} bytes",
                packet.frame_number,
                frame_type_str,
                packet.data.len()
            );

            wav1c::ivf::write_ivf_frame(&mut output, packet.frame_number, &packet.data).unwrap();
        }
    }

    encoder.flush();

    while let Some(packet) = encoder.receive_packet() {
        let frame_type_str = match packet.frame_type {
            wav1c::FrameType::Key => "KEY",
            wav1c::FrameType::Inter => "INTER",
        };
        eprintln!(
            "frame {:>4}  {:>5}  {} bytes",
            packet.frame_number, frame_type_str, packet.data.len()
        );
        wav1c::ivf::write_ivf_frame(&mut output, packet.frame_number, &packet.data).unwrap();
    }

    let mut file = File::create(&cli.output_path).unwrap_or_else(|e| {
        eprintln!("Error creating {}: {}", cli.output_path, e);
        process::exit(1);
    });
    file.write_all(&output).unwrap_or_else(|e| {
        eprintln!("Error writing {}: {}", cli.output_path, e);
        process::exit(1);
    });

    eprintln!();
    if let Some(stats) = encoder.rate_control_stats() {
        eprintln!(
            "Wrote {} bytes to {} ({} frames, target={}kbps, avg_qp={}, buffer={}%, keyint={})",
            output.len(),
            cli.output_path,
            frames.len(),
            stats.target_bitrate / 1000,
            stats.avg_qp,
            stats.buffer_fullness_pct,
            cli.config.keyint
        );
    } else {
        let dq = wav1c::dequant::lookup_dequant(cli.config.base_q_idx);
        eprintln!(
            "Wrote {} bytes to {} ({} frames, q={}, keyint={}, dc_dq={}, ac_dq={})",
            output.len(),
            cli.output_path,
            frames.len(),
            cli.config.base_q_idx,
            cli.config.keyint,
            dq.dc,
            dq.ac
        );
    }
}