use std::{
    path::{Path, PathBuf},
    sync::{Arc, atomic::AtomicBool},
    time::Duration,
};

use dir_lock::DirLock;
use eyre::{Context, Result, eyre};
use indicatif::{ProgressBar, ProgressStyle};
use tokio::{fs, task::JoinSet, time::Instant};
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, instrument};

use crate::{
    error::Error,
    qemu::{
        QemuLaunchOpts, convert_ovmf_uefi_variables, create_overlay_image, extract_kernel,
        launch_qemu,
    },
    ssh::{
        CommandOutput, PersistedSshKeypair, SshLaunchOpts, connect_ssh_for_command,
        connect_ssh_for_warmup, create_ssh_connection, ensure_ssh_key,
    },
    types::{BindMount, EnvVar, Interactive, Io, OsTypeOrImagePath, PmemMount, PublishPort, Pull},
    utils::{ExecutablePaths, VmexecDirs, create_free_cid, find_required_tools, lock_state_dir},
    vm_images::{VmImage, resolve_image},
    vms::{
        SESSION_LOG_FILENAME, VM_CONFIG_FILENAME, VmConfig, VmId, VmReservation, VmState, get_vm,
    },
};

#[derive(Debug, Clone, Default)]
pub(crate) struct CancellationTokens {
    pub qemu: CancellationToken,
    pub ssh: CancellationToken,
}

/// Everything that [`prepare_image`] produced and that a subsequent run needs
#[derive(Debug)]
pub(crate) struct PreparedImage {
    pub overlay_image_path: PathBuf,
    pub cid: u32,
    pub ovmf_vars: PathBuf,
}

/// Knobs for the warmup boot performed by [`prepare_image`]
#[derive(Debug)]
pub struct PrepareOpts {
    pub ssh_timeout: Duration,
    pub memory: u64,
    pub disable_kvm: bool,
    pub show_vm_window: bool,
}

/// Prepare an image for running
///
/// This performs the whole preparation dance for an image: SSH keypair, free CID, OVMF UEFI vars
/// and finally [`warmup`] (kernel extraction, overlay creation and the warmup boot).
///
/// The warmup boot doesn't use volumes, published ports or pmems so those are not configurable
/// here.
#[instrument(skip_all)]
pub(crate) async fn prepare_image(
    dirs: &VmexecDirs,
    lock: &DirLock,
    tool_paths: ExecutablePaths,
    vm_dir: &Path,
    vm_image: VmImage,
    ssh_keypair: PersistedSshKeypair,
    opts: PrepareOpts,
) -> Result<PreparedImage> {
    // We need a free CID for host-guest communication via vsock.
    let cid = create_free_cid(lock, &dirs.vms_dir, vm_dir)?;

    let ovmf_vars_system_path = Path::new("/usr/share/edk2/x64/OVMF_VARS.4m.fd");
    let ovmf_vars = convert_ovmf_uefi_variables(vm_dir, ovmf_vars_system_path).await?;

    let qemu_launch_opts = QemuLaunchOpts {
        volumes: vec![],
        pmems: vec![],
        published_ports: vec![],
        vm_image,
        ovmf_uefi_vars_path: ovmf_vars.clone(),
        show_vm_window: opts.show_vm_window,
        pubkey: ssh_keypair.pubkey_str,
        cid,
        is_warmup: true,
        disable_kvm: opts.disable_kvm,
        memory: opts.memory,
    };

    let ssh_launch_opts = SshLaunchOpts {
        timeout: opts.ssh_timeout,
        tty: false,
        interactive: Interactive::Never,
        env_vars: vec![],
        workdir: None,
        args: vec![],
        privkey: ssh_keypair.privkey_str,
        cid,
        stdout: Io::Inherit,
        stderr: Io::Inherit,
    };

    let overlay_image_path = warmup(vm_dir, tool_paths, qemu_launch_opts, ssh_launch_opts).await?;

    Ok(PreparedImage {
        overlay_image_path,
        cid,
        ovmf_vars,
    })
}

/// Make sure the VM is warmed up
///
/// This will:
/// 1. Extract the kernel and initrd from the image's boot partition so we can boot more quickly
///    without going through the bootloader
/// 2. Create an overlay image with the supplied image as a backing file
/// 3. Spin up the supplied VM
/// 4. Wait until it is connectable via SSH
/// 5. Wait until a readiness condition is met
/// 6. Stop it
#[instrument(skip(tool_paths, qemu_launch_opts, ssh_launch_opts))]
pub(crate) async fn warmup(
    run_dir: &Path,
    tool_paths: ExecutablePaths,
    qemu_launch_opts: QemuLaunchOpts,
    ssh_launch_opts: SshLaunchOpts,
) -> Result<PathBuf> {
    // Before doing anything, first determine whether the image has already been warmed up.
    // We assume that if an overlay image and its extracted kernel exist and will exit early so we
    // don't do any duplicate work.
    let overlay_image_path = qemu_launch_opts.vm_image.overlay_image();
    if overlay_image_path.exists() {
        debug!("Found existing overlay image at {overlay_image_path:?}");

        if qemu_launch_opts.vm_image.kernel_path.exists() {
            debug!("Found extracted kernel, nothing to do for warmup");
            return Ok(overlay_image_path);
        }
    }

    let progress = ProgressBar::new_spinner();
    progress.set_style(
        ProgressStyle::with_template("{spinner:.blue} {msg}")?.tick_strings(&[
            "▹▹▹▹▹",
            "▸▹▹▹▹",
            "▹▸▹▹▹",
            "▹▹▸▹▹",
            "▹▹▹▸▹",
            "▹▹▹▹▸",
            "▪▪▪▪▪",
        ]),
    );
    progress.enable_steady_tick(Duration::from_millis(100));
    progress.set_message("Warming up image");

    progress.println("Extracting kernel image");
    extract_kernel(&tool_paths.virt_copy_out_path, &qemu_launch_opts.vm_image).await?;

    info!("No existing overlay image found, creating...");

    progress.println("Creating overlay image");
    create_overlay_image(&qemu_launch_opts.vm_image.image_path, &overlay_image_path).await?;

    // Create a new launch struct where we use the overlay image instead of the source image.
    let vm_image = VmImage {
        image_path: overlay_image_path.clone(),
        ..qemu_launch_opts.vm_image
    };
    let qemu_launch_opts = QemuLaunchOpts {
        vm_image,
        ..qemu_launch_opts
    };

    // In case of the warmup, we'll poweroff the VM at some point. From that point on, QEMU is
    // expected to exit. If QEMU exit at any other point, it's probably some kind of error.
    // These channels allow us to signal that.
    let qemu_should_exit = Arc::new(AtomicBool::new(false));

    progress.println("Running virtual machine");
    let cancellatation_tokens = CancellationTokens::default();
    let mut joinset = JoinSet::new();
    joinset.spawn({
        let qemu_should_exit = qemu_should_exit.clone();
        let run_dir = run_dir.to_owned();
        let cancellatation_tokens_ = cancellatation_tokens.clone();
        async move {
            launch_qemu(
                cancellatation_tokens_,
                qemu_should_exit,
                run_dir.as_path(),
                None,
                tool_paths,
                qemu_launch_opts,
            )
            .await
        }
    });
    connect_ssh_for_warmup(qemu_should_exit, ssh_launch_opts)
        .await
        .wrap_err("SSH connection error")?;

    while let Some(res) = joinset.join_next().await {
        res??
    }

    progress.finish_with_message("Image is warmed up");

    Ok(overlay_image_path)
}

/// Warm up an image
///
/// The warmup boot needs a VM dir just like a regular run does and so this creates a throwaway one
/// and removes it again afterwards whether or not the warmup succeeded.
///
/// NOTE: Locks the state dir.
#[instrument(skip(dirs, vm_image, opts))]
pub async fn prepare_pulled_image(
    dirs: &VmexecDirs,
    vm_image: VmImage,
    opts: PrepareOpts,
) -> Result<(), Error> {
    // The warmup writes the overlay image into the cache which could race a concurrent `run` of the
    // same image.
    let lock = lock_state_dir(dirs)?;

    let tool_paths = find_required_tools()?;
    let ssh_keypair = ensure_ssh_key(&dirs.secrets_dir)?;

    let vmid = VmId::default();
    let vm_dir = Path::new(&dirs.vms_dir).join(vmid.as_str());
    fs::create_dir(&vm_dir)
        .await
        .wrap_err(format!("Couldn't make temp dir in {:?}", dirs.vms_dir))?;

    let prepare_result = prepare_image(
        dirs,
        &lock,
        tool_paths,
        &vm_dir,
        vm_image,
        ssh_keypair,
        opts,
    )
    .await;

    // Clean up the throwaway VM dir whether or not the preparation worked.
    fs::remove_dir_all(&vm_dir)
        .await
        .wrap_err(format!("Couldn't remove temp dir {vm_dir:?}"))?;
    prepare_result?;

    Ok(())
}

/// Keeps track of whether a VM is reachable over SSH
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Readiness {
    /// The VM is up and accepting SSH connections.
    Ready,

    /// The VM is alive but isn't accepting SSH connections.
    NotReady,

    /// The VM has already exited.
    Exited,
}

/// Check whether a VM is ready to accept SSH connections
///
/// If given a `timeout`, this polls until the VM reports as running and then gives SSH that long to
/// connect.
/// If `timeout` is not provided, it checks the current state and returns immediately.
#[instrument(skip(dirs))]
pub async fn check_ready(
    dirs: &VmexecDirs,
    vmid: &VmId,
    timeout: Option<Duration>,
) -> Result<Readiness, Error> {
    let ssh_keypair = ensure_ssh_key(&dirs.secrets_dir)?;

    // If the VM has already exited, we should exit early.
    let mut vm = get_vm(&dirs.vms_dir, vmid)?;
    if vm.state() == VmState::Exited {
        return Ok(Readiness::Exited);
    }

    if let Some(timeout) = timeout {
        let now = Instant::now();
        while now.elapsed() < timeout {
            vm = get_vm(&dirs.vms_dir, vmid)?;
            if vm.state() == VmState::Running {
                break;
            }
            tokio::time::sleep(Duration::from_secs(1)).await;
        }
    }

    let cid = vm
        .cid_pids
        .cid
        .ok_or_else(|| Error::Other(eyre!("cid didn't exist")))?;

    let ssh_launch_opts = SshLaunchOpts {
        timeout: timeout.unwrap_or(Duration::from_secs(0)),
        tty: false,
        interactive: Interactive::Never,
        env_vars: vec![],
        workdir: None,
        args: vec!["echo".to_string(), "ok".to_string()],
        privkey: ssh_keypair.privkey_str,
        cid,
        stdout: Io::Inherit,
        stderr: Io::Inherit,
    };

    if create_ssh_connection(&ssh_launch_opts).await.is_ok() {
        Ok(Readiness::Ready)
    } else {
        Ok(Readiness::NotReady)
    }
}

#[derive(Debug)]
pub struct ExecOpts {
    pub ssh_timeout: Duration,
    pub tty: bool,
    pub interactive: Interactive,
    pub env_vars: Vec<EnvVar>,
    pub workdir: Option<PathBuf>,
    pub args: Vec<String>,
    pub stdout: Io,
    pub stderr: Io,
}

/// Run a command in an already running VM
///
/// Returns the exit code of the command inside the VM if it reported one.
#[instrument(skip(dirs, exec_opts))]
pub async fn exec_in_running_vm(
    dirs: &VmexecDirs,
    vmid: &VmId,
    exec_opts: ExecOpts,
) -> Result<Option<CommandOutput>, Error> {
    let ssh_keypair = ensure_ssh_key(&dirs.secrets_dir)?;

    let vm = get_vm(&dirs.vms_dir, vmid)?;
    let cid = vm
        .cid_pids
        .cid
        .ok_or_else(|| Error::Other(eyre!("cid didn't exist")))?;
    let session_log_path = vm.vm_dir.join(SESSION_LOG_FILENAME);

    let ssh_launch_opts = SshLaunchOpts {
        timeout: exec_opts.ssh_timeout,
        tty: exec_opts.tty,
        interactive: exec_opts.interactive,
        env_vars: exec_opts.env_vars,
        workdir: exec_opts.workdir,
        args: exec_opts.args,
        privkey: ssh_keypair.privkey_str,
        cid,
        stdout: exec_opts.stdout,
        stderr: exec_opts.stderr,
    };

    Ok(connect_ssh_for_command(None, ssh_launch_opts, session_log_path).await?)
}

/// Knobs for running a command in a new VM via [`prepare_and_run`]
#[derive(Debug)]
pub struct RunOptions {
    pub image_source: OsTypeOrImagePath,
    pub pull: Pull,
    pub ssh_timeout: Duration,
    pub memory: u64,
    pub disable_kvm: bool,
    pub show_vm_window: bool,
    pub tty: bool,
    pub interactive: Interactive,
    pub env_vars: Vec<EnvVar>,
    pub workdir: Option<PathBuf>,
    pub args: Vec<String>,
    pub volumes: Vec<BindMount>,
    pub pmems: Vec<PmemMount>,
    pub published_ports: Vec<PublishPort>,
    pub stdout: Io,
    pub stderr: Io,
    pub rm: bool,
}

/// Prepare an image and run a user-supplied command in a new VM
///
/// This does a lot of heavy lifting:
/// resolve the image, warm it up, and boot the actual VM with the requested command.
#[instrument(skip(dirs, reservation, tool_paths, opts))]
pub async fn prepare_and_run(
    dirs: &VmexecDirs,
    reservation: VmReservation,
    tool_paths: ExecutablePaths,
    opts: RunOptions,
) -> Result<Option<CommandOutput>, Error> {
    let VmReservation {
        vmid,
        vm_dir,
        state_dir_lock: lock,
    } = reservation;
    let vm_dir = vm_dir.as_path();

    let vm_image = match opts.image_source {
        OsTypeOrImagePath::OsType(os_type) => {
            resolve_image(&dirs.images_dir, os_type, opts.pull).await?
        }
        OsTypeOrImagePath::ImagePath(image_path) => VmImage::generic(&image_path)?,
    };

    let ssh_keypair = ensure_ssh_key(&dirs.secrets_dir)?;

    let prepared = prepare_image(
        dirs,
        &lock,
        tool_paths.clone(),
        vm_dir,
        vm_image.clone(),
        ssh_keypair.clone(),
        PrepareOpts {
            ssh_timeout: opts.ssh_timeout,
            memory: opts.memory,
            disable_kvm: opts.disable_kvm,
            show_vm_window: opts.show_vm_window,
        },
    )
    .await?;

    // We launch QEMU from the overlay image instead of the source image.
    let vm_image = VmImage {
        image_path: prepared.overlay_image_path,
        ..vm_image
    };
    let qemu_launch_opts = QemuLaunchOpts {
        volumes: opts.volumes,
        pmems: opts.pmems,
        published_ports: opts.published_ports,
        vm_image,
        ovmf_uefi_vars_path: prepared.ovmf_vars,
        show_vm_window: opts.show_vm_window,
        pubkey: ssh_keypair.pubkey_str,
        cid: prepared.cid,
        is_warmup: false,
        disable_kvm: opts.disable_kvm,
        memory: opts.memory,
    };

    let ssh_launch_opts = SshLaunchOpts {
        timeout: opts.ssh_timeout,
        tty: opts.tty,
        interactive: opts.interactive,
        env_vars: opts.env_vars,
        workdir: opts.workdir,
        args: opts.args,
        privkey: ssh_keypair.privkey_str,
        cid: prepared.cid,
        stdout: opts.stdout,
        stderr: opts.stderr,
    };

    debug!("SSH command for manual debugging:");
    debug!(
        "ssh root@vsock/{cid} -i {privkey_path:?}",
        cid = prepared.cid,
        privkey_path = ssh_keypair.privkey_path
    );

    let output = run_in_new_vm(
        vmid,
        vm_dir,
        lock,
        tool_paths,
        qemu_launch_opts,
        ssh_launch_opts,
        opts.rm,
    )
    .await?;

    // Clean up after running in case this has been requested.
    if opts.rm {
        debug!("Removing VM dir at {vm_dir:?}");
        fs::remove_dir_all(&vm_dir)
            .await
            .wrap_err(format!("Couldn't remove VM dir {vm_dir:?}"))?;
    }

    Ok(output)
}

/// Run a user-supplied command in a throw-away VM
#[instrument(skip(vmid, lock, tool_paths, qemu_launch_opts, ssh_launch_opts))]
pub(crate) async fn run_in_new_vm(
    vmid: VmId,
    vm_dir: &Path,
    lock: DirLock,
    tool_paths: ExecutablePaths,
    qemu_launch_opts: QemuLaunchOpts,
    ssh_launch_opts: SshLaunchOpts,
    rm: bool,
) -> Result<Option<CommandOutput>> {
    let cancellatation_tokens = CancellationTokens::default();
    let mut joinset = JoinSet::new();
    joinset.spawn({
        let vm_data_dir = vm_dir.to_owned();
        let cancellatation_tokens_ = cancellatation_tokens.clone();
        let qemu_launch_opts_ = qemu_launch_opts.clone();

        async move {
            launch_qemu(
                cancellatation_tokens_,
                Arc::new(AtomicBool::new(false)),
                vm_data_dir.as_path(),
                Some(lock),
                tool_paths,
                qemu_launch_opts_,
            )
            .await?;
            Ok(None)
        }
    });
    joinset.spawn({
        let cancellatation_tokens_ = cancellatation_tokens.clone();
        let ssh_launch_opts_ = ssh_launch_opts.clone();
        let session_log_path = vm_dir.join(SESSION_LOG_FILENAME);

        async move {
            connect_ssh_for_command(
                Some(cancellatation_tokens_),
                ssh_launch_opts_,
                session_log_path,
            )
            .await
        }
    });

    let vm_config = VmConfig::new(vmid, qemu_launch_opts, ssh_launch_opts, rm);
    // Write out launch options
    fs::write(
        vm_dir.join(VM_CONFIG_FILENAME),
        serde_json::to_string_pretty(&vm_config)?,
    )
    .await?;

    let mut output = None;
    while let Some(res) = joinset.join_next().await {
        // Workaround to make sure we only return an exit code from SSH.
        // The qemu task will always return None.
        if let Some(actual_output) = res?? {
            output = Some(actual_output);
        };
    }

    Ok(output)
}
