vmexec 0.7.1

use std::{path::PathBuf, str::FromStr, time::Duration};

use clap::{Args, Parser, Subcommand};
use tracing::Level;

use vmexec::types::{
    BindMount, EnvVar, Interactive, Memory, OsType, OsTypeOrImagePath, PmemMount, PublishPort,
    Pull, Tty,
};
use vmexec::utils::VmexecDirs;
use vmexec::vms::{VmId, VmState, get_vms};

/// Parse a string into a Duration
fn parse_seconds_to_duration(src: &str) -> Result<Duration, String> {
    let sec_int = src
        .parse()
        .map_err(|_e| format!("Failed to parse '{src}' as an integer"))?;
    Ok(Duration::from_secs(sec_int))
}

/// Parse and validate that a given `vmid` exists
fn parse_existing_vmid(vmid: &str) -> Result<VmId, String> {
    let dirs = VmexecDirs::new().unwrap();

    let vms = get_vms(&dirs.vms_dir).unwrap();

    let vmid = VmId::new(vmid.to_string())?;
    if vms.contains_key(&vmid) {
        return Ok(vmid);
    }

    Err("No virtual machine with provided ID found".to_string())
}

/// Parse and validate that a given `vmid` exists and is live
fn parse_existing_and_live_vmid(vmid: &str) -> Result<VmId, String> {
    let dirs = VmexecDirs::new().unwrap();

    let vms = get_vms(&dirs.vms_dir).unwrap();

    let vmid = VmId::new(vmid.to_string())?;
    if let Some(vm) = vms.get(&vmid) {
        let vm_is_live = vm.state() == VmState::Running;

        if vm_is_live {
            return Ok(vmid);
        } else {
            return Err(
                "A virtual machine with the provided ID was found but it has already exited"
                    .to_string(),
            );
        }
    }

    Err("No running virtual machine with provided ID found".to_string())
}

#[derive(Debug, Clone, Subcommand)]
pub enum Command {
    /// List virtual machines
    Ps(PsCommand),

    /// Display detailed information on a virtual machine
    Inspect(InspectCommand),

    /// List all available OS images
    Images(ImagesCommand),

    /// Show terminal log of an existing virtual machine
    Logs(LogsCommand),

    /// Stop a virtual machine by sending a SIGTERM signal
    Stop(StopCommand),

    /// Run a command in an existing virtual machine
    Exec(ExecCommand),

    /// Pull an image without running a virtual machine
    Pull(PullCommand),

    /// Run a command in a new virtual machine
    Run(RunCommand),

    /// Check and change KSM status
    ///
    /// Without flags, this prints the current KSM state and some stats.
    Ksm(KsmCommand),

    /// Clean up old virtual machines
    Clean(CleanCommand),

    /// Delete all overlay images from the cache
    Prune(PruneCommand),

    /// Check if a virtual machine is ready to accept SSH connections
    ///
    /// Without --wait, this checks the current state and exits immediately.
    /// With --wait, this polls until the virtual machine is ready or the timeout is reached.
    Ready(ReadyCommand),

    /// Generate completion file for a shell
    Completions { shell: clap_complete::Shell },

    /// Print man page
    Manpage { out_dir: PathBuf },
}

#[derive(Debug, Clone, Args)]
pub struct PsCommand {
    /// Show all the virtual machines, default is only running virtual machines
    #[arg(short, long)]
    pub all: bool,
}

#[derive(Debug, Clone, Args)]
pub struct InspectCommand {
    /// Identifier of a virtual machine
    #[arg(value_parser = parse_existing_vmid)]
    pub vmid: VmId,
}

#[derive(Debug, Clone, Args)]
pub struct ImagesCommand {}

#[derive(Debug, Clone, Args)]
pub struct LogsCommand {
    /// Identifier of a virtual machine
    #[arg(value_parser = parse_existing_vmid)]
    pub vmid: VmId,

    /// Follow log output
    #[arg(short, long)]
    pub follow: bool,
}

#[derive(Debug, Clone, Args)]
pub struct StopCommand {
    /// Identifier of a running virtual machine
    #[arg(value_parser = parse_existing_and_live_vmid)]
    pub vmid: VmId,
}

#[derive(Debug, Clone, Args)]
pub struct ExecCommand {
    /// Set environment variables for the process inside the virtual machine
    ///
    /// Can be provided multiple times.
    ///
    /// Expected format: KEY=VALUE
    #[arg(short, long, value_parser = EnvVar::from_str)]
    pub env: Vec<EnvVar>,

    /// Working directory inside the virtual machine
    ///
    /// The default wokring directory for running binaries within a virtual machine is dependent on
    /// the image. Most commonly, it'll be /root.
    #[arg(short, long)]
    pub workdir: Option<PathBuf>,

    /// SSH connection timeout
    ///
    /// Try for this long (in seconds) to connect to the virtual machine's SSH server.
    #[arg(
        short,
        long,
        default_value = "20",
        value_parser = parse_seconds_to_duration,
    )]
    pub ssh_timeout: Duration,

    /// Make STDIN available to the virtual machine's process
    ///
    /// If 'auto', this will try to read from stdin if it is available, and do nothing when
    /// stdin is not available.
    /// If 'always', this will try to read from stdin and abort when stdin is not available.
    #[arg(short, long, default_value = "auto")]
    pub interactive: Interactive,

    /// Allocate a pseudo-TTY for the virtual machine
    ///
    /// If 'auto', this will be enabled in case vmexec is run from an interactive terminal.
    #[arg(short, long, default_value = "auto")]
    pub tty: Tty,

    /// Identifier of a running virtual machine
    #[arg(value_parser = parse_existing_and_live_vmid)]
    pub vmid: VmId,

    /// Arguments to run in the virtual machine
    pub args: Vec<String>,
}

#[derive(Debug, Clone, Args)]
pub struct PullCommand {
    /// The operating system image to pull
    pub os_type: OsType,

    /// Additionally prepare the image so that the first run is fast
    ///
    /// This extracts the kernel, creates the overlay image and performs the warmup boot that
    /// would otherwise happen on the first run.
    #[arg(long)]
    pub prepare: bool,

    /// Run the warmup virtual machine with TCG instead of KVM
    ///
    /// This would be mostly useful in environments where KVM is not available. It will be a lot
    /// slower but might be acceptable in some contexts.
    #[arg(long, requires = "prepare")]
    pub disable_kvm: bool,

    /// Memory limit for the warmup virtual machine in GBs
    ///
    /// Must be an integer of at least 1.
    ///
    /// If not provided, all available host memory is used.
    #[arg(short, long, requires = "prepare", value_parser = Memory::from_str)]
    pub memory: Option<Memory>,

    /// Show a window with the warmup virtual machine running in it
    ///
    /// This is mostly useful for debugging boot failures.
    #[arg(long, requires = "prepare")]
    pub show_vm_window: bool,

    /// SSH connection timeout during warmup
    ///
    /// Try for this long (in seconds) to connect to the virtual machine's SSH server.
    #[arg(
        short,
        long,
        default_value = "20",
        requires = "prepare",
        value_parser = parse_seconds_to_duration,
    )]
    pub ssh_timeout: Duration,
}

#[derive(Debug, Clone, Args)]
pub struct RunCommand {
    /// Run virtual machine in background and print virtual machine ID
    #[arg(short, long, conflicts_with_all = ["interactive", "tty"])]
    pub detach: bool,

    /// Remove virtual machine runtime data after exit
    #[arg(long)]
    pub rm: bool,

    /// Virtual machine identifier
    ///
    /// If not provided, a random identifier will be generated.
    #[arg(long)]
    pub vmid: Option<VmId>,

    /// Run virtual machine with TCG instead of KVM
    ///
    /// This would be mostly useful in environments where KVM is not available. It will be a lot
    /// slower but might be acceptable in some contexts.
    #[arg(long)]
    pub disable_kvm: bool,

    /// Memory limit for the virtual machine in GBs
    ///
    /// Must be an integer of at least 1.
    ///
    /// If not provided, all available host memory is used.
    #[arg(short, long, value_parser = Memory::from_str)]
    pub memory: Option<Memory>,

    /// Set environment variables for the process inside the virtual machine
    ///
    /// Can be provided multiple times.
    ///
    /// Expected format: KEY=VALUE
    #[arg(short, long, value_parser = EnvVar::from_str)]
    pub env: Vec<EnvVar>,

    /// Working directory inside the virtual machine
    ///
    /// The default wokring directory for running binaries within a virtual machine is dependent on
    /// the image. Most commonly, it'll be /root.
    #[arg(short, long)]
    pub workdir: Option<PathBuf>,

    /// Bind mount a volume into the virtual machine
    ///
    /// Can be provided multiple times.
    ///
    /// Expected format: source:dest[:ro]
    ///
    /// `ro` can optionally be provided to mark the mount as read-only.
    ///
    /// Example: $PWD/src:/mnt:ro
    #[arg(short, long = "volume", value_parser = BindMount::from_str)]
    pub volumes: Vec<BindMount>,

    /// Mount a virtio-pmem device file into the virtual machine
    ///
    /// You might want to do this to bypass the guest page cache. This is important if you're
    /// overprovisioning your host (i.e. giving VMs more combined memory than the host actually
    /// has) and have a write-heavy workload.
    /// For more info, see: https://www.qemu.org/docs/master/system/devices/virtio-pmem.html
    ///
    /// Can be provided multiple times.
    ///
    /// Size is in gigabytes.
    ///
    /// Expected format: dest:<size>
    ///
    /// Example: /var/lib:20
    #[arg(long = "pmem", value_parser = PmemMount::from_str)]
    pub pmems: Vec<PmemMount>,

    /// Publish a port on the virtual machine to the host
    ///
    /// Can be provided multiple times.
    ///
    /// Expected format: [[hostip:][hostport]:]vmport
    ///
    /// `hostip` is optional and if not provided, the port will be bound on all host IPs.
    ///
    /// `hostport` is optional and if not provided, the same value of `vmport` will be used for the
    /// host port.
    ///
    /// Currently only IPv4 is supported for `hostip`.
    #[arg(short, long = "publish", value_parser = PublishPort::from_str)]
    pub published_ports: Vec<PublishPort>,

    /// SSH connection timeout
    ///
    /// Try for this long (in seconds) to connect to the virtual machine's SSH server.
    #[arg(
        short,
        long,
        default_value = "20",
        value_parser = parse_seconds_to_duration,
    )]
    pub ssh_timeout: Duration,

    /// Make STDIN available to the virtual machine's process
    ///
    /// If 'auto', this will try to read from stdin if it is available, and do nothing when
    /// stdin is not available.
    /// If 'always', this will try to read from stdin and abort when stdin is not available.
    #[arg(short, long, default_value = "auto")]
    pub interactive: Interactive,

    /// Allocate a pseudo-TTY for the virtual machine
    ///
    /// If 'auto', this will be enabled in case vmexec is run from an interactive terminal.
    #[arg(short, long, default_value = "auto")]
    pub tty: Tty,

    /// Show a window with the virtual machine running in it
    ///
    /// This is mostly useful for debugging boot failures.
    #[arg(long)]
    pub show_vm_window: bool,

    /// When to pull a new image
    #[arg(long, default_value = "missing")]
    pub pull: Pull,

    /// Either an operating system (e.g. archlinux) or a path to an image
    ///
    /// Possible OS types: archlinux
    #[arg(value_parser = OsTypeOrImagePath::from_str)]
    pub image_source: OsTypeOrImagePath,

    /// Arguments to run in the virtual machine
    pub args: Vec<String>,
}

#[derive(Debug, Clone, Args)]
pub struct KsmEnableDisable {
    /// Persistently enable KSM by writing settings to /etc/tmpfiles.d/ksm.conf
    #[arg(short, long)]
    pub enable: bool,

    /// Persistently disable KSM by deleting /etc/tmpfiles.d/ksm.conf
    #[arg(short, long)]
    pub disable: bool,
}

#[derive(Debug, Clone, Args)]
pub struct KsmCommand {
    #[command(flatten)]
    pub ksm_enable_disable: Option<KsmEnableDisable>,
}

#[derive(Debug, Clone, Args)]
pub struct CleanCommand {}

#[derive(Debug, Clone, Args)]
pub struct PruneCommand {}

#[derive(Debug, Clone, Args)]
pub struct ReadyCommand {
    /// Wait for the virtual machine to become ready
    ///
    /// Without this flag, the command checks the current state and exits immediately.
    #[arg(long)]
    pub wait: bool,

    /// Timeout in seconds when waiting for the virtual machine to become ready
    ///
    /// Only used with --wait flag. Defaults to 60 seconds.
    #[arg(
        long,
        default_value = "60",
        requires = "wait",
        value_parser = parse_seconds_to_duration,
    )]
    pub timeout: Duration,

    /// Identifier of a virtual machine
    #[arg(value_parser = parse_existing_vmid)]
    pub vmid: VmId,
}

/// Run a command in a new virtual machine
#[derive(Debug, Clone, Parser)]
#[command(name = "vmexec", author, about, version)]
pub struct Cli {
    /// Log messages above specified level (error, warn, info, debug, trace)
    #[arg(long, default_value = "warn")]
    pub log_level: Level,

    // Subcommand to run
    #[clap(subcommand)]
    pub command: Command,
}