liblitho 0.2.0

mod cli_cancel;
mod cli_output;

use clap::{Parser, Subcommand};
use cli_cancel::CANCEL_EXIT_CODE;
use cli_output::{CliOutput, OutputMode};
use liblitho::io_backend::{clone_io, flash_io};
use liblitho::progress::is_operation_cancelled;
use std::path::PathBuf;
use std::process::ExitCode;

#[derive(Parser)]
#[command(author, version, about, long_about = None)]
struct Cli {
    /// Output style: terminal progress bar or GUI-friendly line protocol.
    #[arg(short = 'o', long = "output-mode", value_enum, default_value_t = OutputMode::Terminal, global = true)]
    output_mode: OutputMode,

    /// Validate inputs and print the operation that would run, without performing I/O.
    #[arg(long = "dry-run", global = true, default_value_t = false)]
    dry_run: bool,

    /// Path watched for cooperative cancel requests (GUI sidecar / pkexec).
    #[arg(long = "cancel-file", global = true)]
    cancel_file: Option<PathBuf>,

    /// Confirm destructive operations, including automatic unmount of volumes on the target disk.
    #[arg(long = "yes", global = true, default_value_t = false)]
    yes: bool,

    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Read a block device into an image file.
    Clone {
        /// Output image file.
        #[arg(short, long)]
        file: String,

        /// Source block device.
        #[arg(short, long)]
        device: String,

        /// I/O buffer size in bytes.
        #[arg(short, long, default_value_t = 4096)]
        block_size: usize,

        /// Suppress progress output.
        #[arg(short, long, default_value_t = false)]
        silent: bool,
    },
    /// Write an image file to a block device.
    Flash {
        /// Image file to write.
        #[arg(short, long)]
        file: String,

        /// Target block device.
        #[arg(short, long)]
        device: String,

        /// I/O buffer size in bytes.
        #[arg(short, long, default_value_t = 4096)]
        block_size: usize,

        /// Suppress progress output.
        #[arg(short, long, default_value_t = false)]
        silent: bool,

        /// After writing, read the device back and compare SHA-256 checksums.
        #[arg(long = "verify", default_value_t = false)]
        verify: bool,
    },
    /// List storage devices or query one device.
    Query {
        /// Optional device path to query.
        #[arg(short, long)]
        device: Option<String>,
    },
}

fn run(cli: Cli) -> ExitCode {
    let mut out = CliOutput::new(cli.output_mode);

    match cli.command {
        Commands::Clone {
            file,
            device,
            block_size,
            silent,
        } => run_clone(
            &mut out,
            &device,
            &file,
            block_size,
            silent,
            cli.dry_run,
            cli.yes,
            cli.cancel_file.as_deref(),
        ),
        Commands::Flash {
            file,
            device,
            block_size,
            silent,
            verify,
        } => run_flash(
            &mut out,
            &file,
            &device,
            block_size,
            silent,
            verify,
            cli.dry_run,
            cli.yes,
            cli.cancel_file.as_deref(),
        ),
        Commands::Query { device } => run_query(&out, device.as_deref()),
    }
}

fn confirm_volume_dismount(out: &mut CliOutput, device: &str, yes: bool) -> Result<(), ExitCode> {
    let mounts = match liblitho::devices::list_device_mounts(device) {
        Ok(mounts) => mounts,
        Err(e) => {
            out.error(&format!("Could not query volumes on {device}: {e}"));
            return Err(ExitCode::FAILURE);
        }
    };
    if mounts.is_empty() {
        return Ok(());
    }

    let summary: Vec<String> = mounts
        .iter()
        .map(|m| format!("{} on {}", m.source, m.mount_point))
        .collect();

    if !yes {
        out.error(&format!(
            "Disk {device} has mounted volumes ({}). Litho will unmount them automatically before writing. \
             Re-run with --yes to confirm (all data on the disk will be destroyed).",
            summary.join(", ")
        ));
        return Err(ExitCode::FAILURE);
    }

    out.query_status(&format!(
        "Unmounting volumes on {device}: {}",
        summary.join(", ")
    ));
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn run_flash(
    out: &mut CliOutput,
    file: &str,
    device: &str,
    block_size: usize,
    silent: bool,
    verify: bool,
    dry_run: bool,
    yes: bool,
    cancel_file: Option<&std::path::Path>,
) -> ExitCode {
    if !dry_run {
        if let Err(e) = liblitho::devices::validate_device_for_io(device) {
            out.error(&e.to_string());
            return ExitCode::FAILURE;
        }
        if let Err(code) = confirm_volume_dismount(out, device, yes) {
            return code;
        }
        if let Err(e) = liblitho::devices::ensure_device_ready_for_io(device, yes) {
            out.error(&e.to_string());
            return ExitCode::FAILURE;
        }
    }

    if dry_run {
        out.dry_run_ok("flash", file, device, block_size);
        return ExitCode::SUCCESS;
    }

    out.operation_start("Flashing", file, device, block_size);

    let cancel = cli_cancel::prepare_operation_cancel();
    cli_cancel::spawn_cancel_watchers(cancel.clone(), cancel_file.map(PathBuf::from));
    let cancel_ref = Some(cancel.as_ref());
    let result = if silent {
        flash_io::<fn(liblitho::progress::OperationProgress)>(
            file, device, block_size, true, verify, None, cancel_ref,
        )
    } else {
        flash_io(
            file,
            device,
            block_size,
            false,
            verify,
            Some(|event| {
                out.on_progress(&event);
            }),
            cancel_ref,
        )
    };

    out.finish_progress_line();

    match result {
        Ok(()) => {
            out.done_ok("flash");
            ExitCode::SUCCESS
        }
        Err(e) if is_operation_cancelled(&e) => {
            out.cancelled("Flash cancelled - device may be partially written.");
            ExitCode::from(CANCEL_EXIT_CODE)
        }
        Err(e) => {
            out.error(&e.to_string());
            ExitCode::FAILURE
        }
    }
}

#[allow(clippy::too_many_arguments)]
fn run_clone(
    out: &mut CliOutput,
    device: &str,
    file: &str,
    block_size: usize,
    silent: bool,
    dry_run: bool,
    yes: bool,
    cancel_file: Option<&std::path::Path>,
) -> ExitCode {
    if !dry_run {
        if let Err(e) = liblitho::devices::validate_device_for_io(device) {
            out.error(&e.to_string());
            return ExitCode::FAILURE;
        }
        if let Err(code) = confirm_volume_dismount(out, device, yes) {
            return code;
        }
        if let Err(e) = liblitho::devices::ensure_device_ready_for_io(device, yes) {
            out.error(&e.to_string());
            return ExitCode::FAILURE;
        }
    }

    if dry_run {
        out.dry_run_ok("clone", device, file, block_size);
        return ExitCode::SUCCESS;
    }

    out.operation_start("Cloning", device, file, block_size);

    let cancel = cli_cancel::prepare_operation_cancel();
    cli_cancel::spawn_cancel_watchers(cancel.clone(), cancel_file.map(PathBuf::from));
    let cancel_ref = Some(cancel.as_ref());
    let result = if silent {
        clone_io::<fn(liblitho::progress::OperationProgress)>(
            device, file, block_size, true, None, cancel_ref,
        )
    } else {
        clone_io(
            device,
            file,
            block_size,
            false,
            Some(|event| {
                out.on_progress(&event);
            }),
            cancel_ref,
        )
    };

    out.finish_progress_line();

    match result {
        Ok(()) => {
            out.done_ok("clone");
            ExitCode::SUCCESS
        }
        Err(e) if is_operation_cancelled(&e) => {
            out.cancelled("Clone cancelled — incomplete output file removed.");
            ExitCode::from(CANCEL_EXIT_CODE)
        }
        Err(e) => {
            out.error(&e.to_string());
            ExitCode::FAILURE
        }
    }
}

fn run_query(out: &CliOutput, device: Option<&str>) -> ExitCode {
    match device {
        Some(path) => {
            out.query_status(&format!("Querying device: {path}"));
            // Single-device lookup is not implemented yet; list all and let the user filter.
            match liblitho::devices::get_storage_devices() {
                Ok(devices) => {
                    let mut found = false;
                    for dev in devices {
                        if device_path_matches(&dev.device_name, path) {
                            out.query_device(&dev);
                            found = true;
                            break;
                        }
                    }
                    if !found {
                        out.error(&format!("Device not found: {path}"));
                        return ExitCode::FAILURE;
                    }
                    ExitCode::SUCCESS
                }
                Err(e) => {
                    out.error(&e.to_string());
                    ExitCode::FAILURE
                }
            }
        }
        None => match liblitho::devices::get_storage_devices() {
            Ok(devices) => {
                if devices.is_empty() {
                    out.query_status("No storage devices found");
                } else {
                    for dev in devices {
                        out.query_device(&dev);
                    }
                }
                ExitCode::SUCCESS
            }
            Err(e) => {
                out.error(&e.to_string());
                ExitCode::FAILURE
            }
        },
    }
}

fn device_path_matches(device_name: &str, query_path: &str) -> bool {
    liblitho::devices::device_paths_equivalent(device_name, query_path)
}

fn main() -> ExitCode {
    run(Cli::parse())
}