rusta-cli 1.3.0

use std::io::{BufRead, IsTerminal, Write};
use std::time::Duration;

use crate::cli::CreateArgs;
use crate::error::{Error, Result};
use crate::io as rio;
use crate::paths;
use crate::provision;
use crate::registry;
use crate::state;
use crate::tart;

pub fn run(args: CreateArgs) -> Result<u8> {
    validate_name_opt(args.vm.as_deref())?;
    if let Some(gui) = args.gui.as_deref() {
        if provision::display_manager_for(gui).is_none() {
            return Err(Error::msg(format!(
                "--gui accepts: ubuntu-desktop, xubuntu-desktop, lubuntu-desktop, lightdm (got '{gui}')"
            )));
        }
    }

    // The image family to clone: an explicit `--image`, else the first
    // configured image (seeded `ubuntu`). Ignored when `--image-ref` is given.
    let selected_image = match args.image.as_deref() {
        Some(i) => registry::validate_image(i)?,
        None => state::images()
            .into_iter()
            .next()
            .unwrap_or_else(|| state::DEFAULT_IMAGE.to_string()),
    };

    let suggested = format!("{}-{}", selected_image, args.version.replace('.', ""));
    let vm_name = match args.vm.clone() {
        Some(n) => n,
        None => {
            if !std::io::stdin().is_terminal() {
                return Err(Error::msg(format!(
                    "VM name is required for `rusta create` in non-interactive contexts. \
                     Pass a name on the command line (e.g. `rusta create {suggested}`)."
                )));
            }
            let picked = prompt_for_name(
                &suggested,
                &mut std::io::stdin().lock(),
                &mut std::io::stdout(),
            )?;
            validate_name_opt(Some(&picked))?;
            picked
        }
    };

    let variant = if args.gui.is_some() { "desktop" } else { "server" };
    println!();
    let banner = match args.image_ref.as_deref() {
        Some(r) => format!("{r} — {variant} — Tart OCI"),
        None => format!("{selected_image} {} — {variant} — Tart OCI", args.version),
    };
    rio::info(&banner);
    println!();

    if tart::exists(&vm_name)? {
        rio::skip(&format!("VM '{vm_name}' already exists"));
        let mut hint = format!(
            "rusta delete {vm_name} && rusta create {vm_name} --version {}",
            args.version
        );
        if let Some(gui) = args.gui.as_deref() {
            hint.push_str(&format!(" --gui {gui}"));
        }
        if let Some(img) = args.image.as_deref() {
            hint.push_str(&format!(" --image {img}"));
        }
        if let Some(r) = args.image_ref.as_deref() {
            hint.push_str(&format!(" --image-ref {r}"));
        }
        if let Some(src) = args.source.as_deref() {
            hint.push_str(&format!(" --source {src}"));
        }
        rio::info(&format!("To recreate: {hint}"));
        return Ok(0);
    }

    let image = resolve_image(&args, &selected_image)?;
    rio::info(&format!("Cloning {image}..."));
    tart::clone_image(&image, &vm_name)?;
    tart::set_resources(&vm_name, args.cpus, args.memory, args.disk)?;
    let _ = state::set_vm_gui(&vm_name, args.gui.is_some());
    rio::ok(&format!(
        "VM created: {} ({} CPUs, {} GB RAM, {} GB disk)",
        vm_name,
        args.cpus,
        args.memory / 1024,
        args.disk
    ));

    // Generate and persist provisioning script for debuggability.
    let script = provision::generate(&provision::Spec {
        ubuntu_version: &args.version,
        gui: args.gui.as_deref(),
    });
    paths::ensure_dirs().map_err(|e| Error::msg(e.to_string()))?;
    let script_path = paths::provision_script(&vm_name);
    std::fs::write(&script_path, &script).map_err(|e| Error::msg(e.to_string()))?;
    rio::ok(&format!("Provisioning script: {}", script_path.display()));

    // Boot, wait for guest agent, upload + run provisioning, shut down.
    let headless = !args.debug_no_headless;
    if headless {
        rio::info(&format!("Starting VM '{vm_name}' headlessly..."));
    } else {
        rio::info(&format!("Starting VM '{vm_name}' with graphics window (debug)..."));
    }
    let child = tart::run_background(&vm_name, headless)?;
    let pid = child.id();
    let _ = tart::write_pid_file(&vm_name, pid);
    std::mem::forget(child);

    let cleanup = ProcessGuard { pid };

    rio::info("Waiting for guest agent...");
    tart::wait_for_guest_agent(&vm_name, Duration::from_secs(120))?;
    rio::ok("Guest agent is ready");

    rio::info("Uploading provisioning script to guest...");
    tart::exec_with_stdin(
        &vm_name,
        &["bash", "-c", "cat > /tmp/provision.sh && chmod +x /tmp/provision.sh"],
        script.as_bytes(),
    )?;

    rio::info("Running provisioning inside the guest (this may take a while)...");
    tart::exec(&vm_name, &["bash", "/tmp/provision.sh"])?;
    rio::ok("Provisioning complete!");

    rio::info("Shutting down the guest...");
    let _ = tart::exec_quiet(&vm_name, &["sudo", "shutdown", "-h", "now"]);
    let deadline = std::time::Instant::now() + Duration::from_secs(120);
    while std::time::Instant::now() < deadline {
        if !tart::is_running(&vm_name)? {
            break;
        }
        std::thread::sleep(Duration::from_secs(1));
    }
    tart::remove_pid_file(&vm_name);
    rio::ok("VM stopped after provisioning");
    drop(cleanup);

    if args.ssh_copy_keys {
        println!();
        rio::info(&format!("Copying host SSH configuration into '{vm_name}'..."));
        crate::commands::ssh_copy::run(crate::cli::VmOnlyArgs { vm: Some(vm_name.clone()) })?;
    }

    println!();
    rio::ok(&format!("Setup complete: {vm_name}"));
    println!("  Guest user   : {} / {}", args.user, args.password);
    println!("  Start VM     : rusta up {vm_name}");
    println!("  Get IP       : rusta ip {vm_name}");
    println!("  SSH          : rusta ssh {vm_name}");

    Ok(0)
}

/// Resolve the image reference to clone for the selected `image` family.
/// Precedence:
///   1. `--image-ref <ref>` — used verbatim.
///   2. `--source <reg>` — pinned to a single configured source.
///   3. all configured sources in priority order.
///
/// A single candidate is used directly with no registry query (preserves the
/// offline single-source path). With two or more candidates, each source's tags
/// are queried and the first (by priority) advertising `--version` wins;
/// unreachable sources are skipped with a warning.
fn resolve_image(args: &CreateArgs, image: &str) -> Result<String> {
    if let Some(image_ref) = &args.image_ref {
        return Ok(image_ref.clone());
    }
    let candidates = candidate_sources(args)?;
    if candidates.len() == 1 {
        return Ok(candidates[0].image_ref(image, &args.version));
    }
    let results: Vec<(state::Source, std::result::Result<Vec<String>, String>)> = candidates
        .iter()
        .map(|s| (s.clone(), registry::fetch_tags(s, image).map_err(|e| e.message)))
        .collect();
    let pick = registry::pick_image(image, &args.version, &results);
    for label in &pick.warn {
        rio::skip(&format!("source '{label}' unreachable, skipping"));
    }
    match pick.image {
        Some(img) => Ok(img),
        None => Err(Error::msg(pick.err.unwrap_or_else(|| {
            "no configured source provided the requested version".to_string()
        }))),
    }
}

/// Candidate sources for resolution: a single pinned `--source`, or all
/// configured sources in priority order.
fn candidate_sources(args: &CreateArgs) -> Result<Vec<state::Source>> {
    let all = state::sources();
    match &args.source {
        None => Ok(all),
        Some(reg) => {
            let norm = registry::normalize_registry(reg);
            all.into_iter()
                .find(|s| s.registry == norm || s.label() == reg.as_str())
                .map(|s| vec![s])
                .ok_or_else(|| {
                    Error::msg(format!(
                        "source '{reg}' is not configured (run `rusta source add {reg}` or see `rusta source list`)"
                    ))
                })
        }
    }
}

/// Interactively prompt for a VM name, offering `suggested` as the default.
/// An empty line accepts the suggestion; EOF aborts.
pub(crate) fn prompt_for_name<R: BufRead, W: Write>(
    suggested: &str,
    input: &mut R,
    out: &mut W,
) -> Result<String> {
    write!(out, "VM name [{suggested}]: ").map_err(|e| Error::msg(e.to_string()))?;
    out.flush().ok();
    let mut buf = String::new();
    let n = input
        .read_line(&mut buf)
        .map_err(|e| Error::msg(e.to_string()))?;
    if n == 0 {
        return Err(Error::msg("aborted: no VM name provided".to_string()));
    }
    let trimmed = buf.trim();
    if trimmed.is_empty() {
        Ok(suggested.to_string())
    } else {
        Ok(trimmed.to_string())
    }
}

pub(crate) fn validate_name_opt(name: Option<&str>) -> Result<()> {
    let Some(name) = name else { return Ok(()) };
    let first_ok = name
        .chars()
        .next()
        .map(|c| c.is_ascii_alphanumeric())
        .unwrap_or(false);
    let rest_ok = name
        .chars()
        .all(|c| c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-'));
    if !first_ok || !rest_ok {
        return Err(Error::msg(format!(
            "invalid VM name '{name}' (must match ^[a-zA-Z0-9][a-zA-Z0-9._-]*$)"
        )));
    }
    Ok(())
}

struct ProcessGuard {
    pid: u32,
}

impl Drop for ProcessGuard {
    fn drop(&mut self) {
        if tart::pid_alive(self.pid) {
            tart::kill_pid(self.pid);
        }
    }
}

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

    #[test]
    fn names_accept_alnum_dot_underscore_dash() {
        for n in ["a", "ubuntu-2404", "VM.1", "x_y", "abc123"] {
            assert!(validate_name_opt(Some(n)).is_ok(), "should accept {n}");
        }
    }

    #[test]
    fn names_reject_invalid() {
        for n in ["", "-foo", ".bar", "_baz", "has space", "x/y", "x:y", "x@y"] {
            assert!(validate_name_opt(Some(n)).is_err(), "should reject '{n}'");
        }
    }

    #[test]
    fn none_name_is_ok() {
        assert!(validate_name_opt(None).is_ok());
    }

    fn ask(input: &str, suggested: &str) -> Result<String> {
        let mut out = Vec::<u8>::new();
        let mut reader = std::io::Cursor::new(input.as_bytes().to_vec());
        prompt_for_name(suggested, &mut reader, &mut out)
    }

    #[test]
    fn prompt_empty_line_accepts_suggested() {
        assert_eq!(ask("\n", "ubuntu-2404").unwrap(), "ubuntu-2404");
    }

    #[test]
    fn prompt_explicit_name_overrides_suggested() {
        assert_eq!(ask("lab\n", "ubuntu-2404").unwrap(), "lab");
    }

    #[test]
    fn prompt_trims_whitespace() {
        assert_eq!(ask("  lab  \n", "ubuntu-2404").unwrap(), "lab");
    }

    #[test]
    fn prompt_eof_aborts() {
        let err = ask("", "ubuntu-2404").unwrap_err();
        assert!(err.message.contains("aborted"));
    }

    #[test]
    fn prompt_renders_suggested_in_brackets() {
        let mut out = Vec::<u8>::new();
        let mut reader = std::io::Cursor::new(b"\n".to_vec());
        prompt_for_name("ubuntu-2204", &mut reader, &mut out).unwrap();
        let text = String::from_utf8(out).unwrap();
        assert!(text.contains("[ubuntu-2204]"));
    }
}