vmexec 0.7.1

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::time::SystemTime;

use eyre::{Context, OptionExt, Report, Result, bail};
use indicatif::{ProgressBar, ProgressStyle};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use tokio::{fs, io::AsyncWriteExt, task};
use tracing::{debug, info, instrument, trace};
use url::Url;
use walkdir::WalkDir;

use crate::{
    types::{OsType, Pull},
    utils::{VmexecDirs, ensure_directory_async, lock_state_dir},
};

/// Extension of overlay images in the cache dir
pub(crate) const OVERLAY_IMAGE_EXTENSION: &str = "overlay.qcow2";

/// An VM image and auxiliary files
///
/// This consists of a path to the image on disk as well as kernel and initrd file names.
///
/// `initrd` is optional since some distributions (such as Arch Linux) bake all required modules
/// directly into the kernel. This shaves a few hundred milliseconds off the boot so we won't use an
/// initrd if we can avoid it.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VmImage {
    pub image_path: PathBuf,
    pub kernel_path: PathBuf,
    pub initrd_path: Option<PathBuf>,
}

impl VmImage {
    /// Construct a `VmImage` from an `image_path`
    pub fn generic(image_path: &Path) -> Result<Self> {
        let parent = image_path.parent().ok_or_eyre(format!(
            "custom image path {image_path:?} must have a parent directory to store kernel and initramfs"
        ))?;

        Ok(Self {
            image_path: image_path.to_path_buf(),
            kernel_path: parent.join("vmlinuz-linux"),
            initrd_path: Some(parent.join("initramfs-linux.img")),
        })
    }

    /// Construct a `VmImage` for Arch Linux
    pub async fn archlinux(images_dir: &Path, pull: Pull) -> Result<Self> {
        let image_path = ensure_archlinux_image(images_dir, pull).await?;
        Ok(Self {
            image_path: image_path.clone(),
            kernel_path: image_path.parent().unwrap().join("vmlinuz-linux"),
            initrd_path: None,
        })
    }

    /// Get the overlay image path
    pub fn overlay_image(&self) -> PathBuf {
        self.image_path.with_extension(OVERLAY_IMAGE_EXTENSION)
    }
}

/// Download Arch Linux image
pub(crate) async fn download_archlinux_image(
    local_image_path: &Path,
    local_image_checksum_path: &Path,
    arch_boxes_base_url: &Url,
    image_name: &str,
) -> Result<()> {
    let image_url = arch_boxes_base_url.join(&format!("output/{image_name}?job=build:secure"))?;
    let image_checksum_url =
        arch_boxes_base_url.join(&format!("output/{image_name}.SHA256?job=build:secure"))?;

    trace!("Getting Arch Linux image checksum from '{image_checksum_url}'");

    let mut image_checksum = reqwest::get(image_checksum_url)
        .await?
        .error_for_status()?
        .bytes()
        .await?;

    fs::write(&local_image_checksum_path, &image_checksum)
        .await
        .wrap_err(format!(
            "Can't write checksum file to {local_image_checksum_path:?}"
        ))?;

    let mut local_image_file = fs::OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .truncate(true)
        .open(&local_image_path)
        .await?;

    trace!("Resolving Arch Linux image at {image_url}");

    let mut image_resp = reqwest::get(image_url.clone()).await?.error_for_status()?;
    let image_size = image_resp
        .content_length()
        .ok_or_eyre("Couldn't get image size")?;

    debug!("Resolved as {} with {} bytes", image_resp.url(), image_size);

    let progress = ProgressBar::new(image_size);
    progress.set_style(
        ProgressStyle::default_bar()
            .template("{msg}\n{spinner:.green} [{bar:40.green/black}] {bytes:>11.green}/{total_bytes:<11.green} {bytes_per_sec:>13.red} eta {eta:.blue}")?
            .progress_chars("█▉▊▋▌▍▎▏  "),
    );
    progress.set_message(format!("Downloading {}", image_resp.url()));

    let mut hasher = Sha256::new();
    while let Some(chunk) = image_resp.chunk().await? {
        local_image_file.write_all(&chunk).await?;
        hasher.update(&chunk);
        progress.inc(chunk.len() as u64);
    }
    progress.finish_with_message("Download complete!");

    info!("Checking file checksum");

    // The hash file we downloaded is in the format "hash filename" so we'll have to cut off the
    // first part to get just the hash.
    image_checksum.truncate(64);
    let image_checksum_raw = hex::decode(image_checksum)?;

    let computed_checksum = hasher.finalize();
    if image_checksum_raw != computed_checksum.as_slice() {
        bail!(
            "Checksum mismatch on {local_image_path:?}, maybe the file got corrupted somehow. Try deleting it and retrying."
        );
    }

    Ok(())
}

/// Retrieve the latest locally available Arch Linux image
pub(crate) async fn get_latest_local_archlinux_image(
    distro_image_dir: &Path,
) -> Result<Option<PathBuf>> {
    let mut images = task::spawn_blocking({
        let distro_image_dir_ = distro_image_dir.to_owned();
        move || {
            let mut images_ = vec![];

            // First we'll gather a list of all local images.
            for entry in WalkDir::new(distro_image_dir_) {
                let entry = entry?;
                let filename = entry.file_name();
                if entry.file_type().is_file()
                    && filename.to_string_lossy().starts_with("Arch-Linux")
                    && entry.path().extension().unwrap_or_default() == "qcow2"
                {
                    images_.push(entry.path().to_owned());
                }
            }
            Ok::<Vec<PathBuf>, Report>(images_)
        }
    })
    .await??;

    // Sort the images by name as the Arch Linux images contain the ISO date so that the latest one
    // will be the one at the very end.
    images.sort();

    let latest_image = images.last().cloned();
    Ok(latest_image)
}

/// A locally cached VM image
#[derive(Debug)]
pub struct CachedImage {
    pub version: String,
    pub size: u64,
    pub prepared: bool,
    pub created_at: SystemTime,
}

/// Get all locally cached VM images, grouped by OS
///
/// Within each OS, the versions are sorted with the newest first.
#[instrument]
pub fn get_images(images_dir: &Path) -> Result<BTreeMap<String, Vec<CachedImage>>> {
    let mut images = BTreeMap::new();

    // The cache dir looks like this: .cache/vmexec/images/<os>/<version> so we'll first iterate the
    // OSs and then the versions and find a qcow2 for every version.
    for distro_entry in std::fs::read_dir(images_dir)? {
        let distro_dir = distro_entry?.path();
        if !distro_dir.is_dir() {
            continue;
        }
        let os = distro_dir
            .file_name()
            .unwrap()
            .to_string_lossy()
            .to_string();

        let mut versions = Vec::new();
        for version_entry in std::fs::read_dir(&distro_dir)? {
            let version_dir = version_entry?.path();
            if !version_dir.is_dir() {
                continue;
            }
            let version = version_dir
                .file_name()
                .unwrap()
                .to_string_lossy()
                .to_string();

            for file_entry in std::fs::read_dir(&version_dir)? {
                let path = file_entry?.path();
                // Overlay images also have a `qcow2` extension so we have to explicitly
                // exclude them.
                if path.is_file()
                    && path.extension().unwrap_or_default() == "qcow2"
                    && let Some(name) = path.file_name()
                    && !name
                        .to_string_lossy()
                        .ends_with(&format!(".{OVERLAY_IMAGE_EXTENSION}"))
                {
                    let metadata = path.metadata()?;

                    // An image counts as prepared if both its overlay and its extracted kernel
                    // exist.
                    let overlay_exists = path.with_extension(OVERLAY_IMAGE_EXTENSION).exists();
                    let kernel_exists = version_dir.join("vmlinuz-linux").exists();

                    versions.push(CachedImage {
                        version: version.clone(),
                        size: metadata.len(),
                        prepared: overlay_exists && kernel_exists,
                        created_at: metadata.modified()?,
                    });
                }
            }
        }

        // Newest versions first, like `docker images`.
        versions.sort_by(|a, b| b.version.cmp(&a.version));
        images.insert(os, versions);
    }

    Ok(images)
}

/// Delete every overlay image found in the cache dir
///
/// Returns the paths that were removed.
///
/// NOTE: Locks the state dir.
#[instrument]
pub fn prune_overlays(dirs: &VmexecDirs) -> Result<Vec<PathBuf>> {
    // Overlay images in the cache dir are created during image warmup and by pull --prepare both of
    // which hold the lock.
    // Without it, it would be possible to delete an overlay image in the tiny window between warmup
    // and QEMU opening it which then fails the launch.
    let _lock = lock_state_dir(dirs)?;

    let mut deleted_files = Vec::new();

    for entry in WalkDir::new(&dirs.cache_dir) {
        let entry = entry?;
        let path = entry.path();
        if path.is_file()
            && let Some(name) = path.file_name()
            && name
                .to_string_lossy()
                .ends_with(&format!(".{OVERLAY_IMAGE_EXTENSION}"))
        {
            debug!("Removing overlay image {path:?}");
            std::fs::remove_file(path)?;
            deleted_files.push(path.to_path_buf());
        }
    }

    Ok(deleted_files)
}

/// Download an OS image into the cache if it isn't already there
///
/// NOTE: Locks the state dir.
#[instrument(skip(dirs))]
pub async fn pull_image(dirs: &VmexecDirs, os_type: OsType, pull: Pull) -> Result<VmImage> {
    let _lock = lock_state_dir(dirs)?;
    resolve_image(&dirs.images_dir, os_type, pull).await
}

#[instrument]
pub(crate) async fn resolve_image(
    images_dir: &Path,
    os_type: OsType,
    pull: Pull,
) -> Result<VmImage> {
    match os_type {
        OsType::Archlinux => VmImage::archlinux(images_dir, pull).await,
    }
}

/// Download the most recent Arch Linux image
#[instrument]
pub(crate) async fn ensure_archlinux_image(images_dir: &Path, pull: Pull) -> Result<PathBuf> {
    let distro_image_dir = images_dir.join("archlinux");
    ensure_directory_async("distro image", &distro_image_dir).await?;

    let latest_local_image = get_latest_local_archlinux_image(&distro_image_dir).await?;

    match pull {
        // If --pull missing was provided and if there is already a local image (no matter the age),
        // we'll just return that. If there's no image we'll do nothing which will cause it to be
        // downloaded later in the function.
        Pull::Missing => {
            if let Some(latest) = latest_local_image {
                info!(
                    "Found local image {latest:?} and \"--pull missing\" was provided so this is the image we're using"
                );
                return Ok(latest);
            }
        }
        // If --pull missing was provided and if there is already a local image (no matter the age),
        // we'll just return that. If there's no image we'll error out.
        Pull::Never => {
            if let Some(latest) = latest_local_image {
                info!(
                    "Found local image {latest:?} and \"--pull never\" was provided so this is the image we're using"
                );
                return Ok(latest);
            } else {
                bail!("No local image found and `--pull never` selected, bailing");
            }
        }
        Pull::Newer => {
            if let Some(latest) = latest_local_image {
                info!(
                    "Found local image {latest:?} but there might be a newer image so we're checking"
                );
            }
        }
    }

    let arch_boxes_base_url = Url::parse(
        "https://gitlab.archlinux.org/archlinux/arch-boxes/-/jobs/artifacts/master/raw/",
    )?;
    let build_version_url = arch_boxes_base_url.join("build.env?job=build:secure")?;

    // Figure out what the current version is.
    let build_version = reqwest::get(build_version_url)
        .await?
        .error_for_status()?
        .text()
        .await?;
    let build_version_line = build_version
        .lines()
        .next()
        .ok_or_eyre("No line break in output")?;
    let build_version = build_version_line
        .split('=')
        .next_back()
        .ok_or_eyre(format!(
            "BUILD_VERSION line not in expected format: {build_version_line}"
        ))?;

    let image_name = format!("Arch-Linux-x86_64-libvirt-executor-{build_version}.qcow2");

    debug!("Latest remote image is: {image_name}");

    // The file will be downloaded to this path.
    let image_dir = distro_image_dir.join(build_version);
    ensure_directory_async("image dir", &image_dir).await?;

    let local_image_path = image_dir.join(&image_name);
    let image_ext = local_image_path
        .extension()
        .ok_or_eyre("Somehow the image '{local_image_path:?}' didn't have a file extension")?
        .to_str()
        .ok_or_eyre("File extension in '{image_ext}' isn't ASCII")?;
    let local_image_checksum_path = local_image_path.with_extension(format!("{image_ext}.SHA256"));

    // First check if the files even need downloading or whether we already have them.
    if local_image_path.exists() && local_image_checksum_path.exists() {
        debug!("Found pre-existing files for image, skipping download");
        return Ok(local_image_path);
    } else if pull == Pull::Newer {
        debug!("Didn't find requested image locally, proceeding to download");
    }

    // Since we got to this point we'll need to download the image here.
    download_archlinux_image(
        &local_image_path,
        &local_image_checksum_path,
        &arch_boxes_base_url,
        &image_name,
    )
    .await?;

    Ok(local_image_path)
}