//! Build: mount a podman image, inject /slim/ scripts, and pack the rootfs
//! as a gzipped cpio initrd or qcow2 disk in the registry.

use anyhow::{Context, Result, bail};
use clap::{Args, ValueEnum};
use flate2::write::GzEncoder;
use std::fs::{self, File};
use std::io;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use tempfile::NamedTempFile;

use crate::command::cmd;
use crate::inject;
use crate::registry::registry_dir;

#[derive(Args, Debug)]
pub struct BuildCmd {
    kind: ImageKind,

    /// Image tag / registry subdir name
    name: String,

    /// Override the command to exec in the VM. Inferred from the image's
    /// CMD/ENTRYPOINT if not provided.
    #[clap(long)]
    cmd: Option<String>,
}

#[derive(ValueEnum, Clone, Debug)]
enum ImageKind {
    Qcow2,
    Initrd,
}

pub(crate) fn build(
    BuildCmd {
        kind,
        name,
        cmd: cmd_override,
    }: BuildCmd,
) -> Result<()> {
    let image = &name;
    let container = format!("slim-build-{}", image.replace(':', "-"));

    let mount_path = mount_container(image, &container)?;
    println!("Mounted at: {}", mount_path.display());

    let result = build_inner(&kind, image, &mount_path, cmd_override);

    let unmounted = cmd(&[
        "podman",
        "unshare",
        "--",
        "podman",
        "container",
        "unmount",
        &container,
    ])
    .is_ok();
    let removed = cmd(&["podman", "rm", &container]).is_ok();
    if unmounted && removed {
        println!("Unmounted and removed container.");
    } else {
        eprintln!("warning: failed to clean up container '{container}'");
    }

    result
}

fn build_inner(
    kind: &ImageKind,
    image: &str,
    mount_path: &Path,
    cmd: Option<String>,
) -> Result<()> {
    let config = inject::inspect_config(image)?;
    let command = cmd.unwrap_or_else(|| inject::infer_command(&config));
    let exec_script =
        inject::build_exec_script(&command, &config.env, config.working_dir.as_deref());
    inject::inject(mount_path, &exec_script, config.user.as_deref())?;
    println!("Injected /slim/ (init + exec)");

    match kind {
        ImageKind::Initrd => build_initrd(image, mount_path),
        ImageKind::Qcow2 => build_qcow2(image, mount_path),
    }
}

fn mount_container(image: &str, container: &str) -> Result<PathBuf> {
    cmd(&["podman", "create", "--name", container, image, "/bin/true"])?;
    let mount_path = cmd(&[
        "podman",
        "unshare",
        "--",
        "podman",
        "container",
        "mount",
        container,
    ])?;
    Ok(PathBuf::from(mount_path.trim()))
}

/// Copy a directory onto a new raw disk image with EXT4.
fn to_raw_ext4(dir: &Path, tmp_dir: &Path) -> Result<NamedTempFile> {
    let dir = dir.to_str().context("Invalid UTF-8")?;
    let du_out = cmd(&["podman", "unshare", "--", "du", "-sk", dir])?;

    let used_kb: u64 = du_out
        .split_whitespace()
        .next()
        .context("du produced no output")?
        .parse()
        .context("failed to parse du output")?;
    let used = used_kb * 1024;
    let gb = 1024 * 1024 * 1024;
    let size = used
        + (used / 10) // Add 10%
        // TODO: make configurable
        + 64 * gb; // Add some spare capacity for activities
    let size = size.to_string();

    let raw_file = NamedTempFile::new_in(tmp_dir)?;
    let raw_path = raw_file.path().to_str().context("Invalid UTF-8")?;
    cmd(&["podman", "unshare", "--", "truncate", "-s", &size, raw_path])?;
    cmd(&[
        "podman",
        "unshare",
        "--",
        "mkfs.ext4",
        "-F",
        "-d",
        dir,
        raw_path,
    ])?;
    Ok(raw_file)
}

/// Copy a directory onto a new qcow2 disk image with EXT4.
fn to_qcow2_ext4(mount_path: &Path, tmp_dir: &Path) -> Result<NamedTempFile> {
    let raw = to_raw_ext4(mount_path, tmp_dir)?;
    let qcow2 = NamedTempFile::new_in(tmp_dir)?;
    let raw_path = raw.path().to_str().context("Invalid UTF-8")?;
    let qcow2_path = qcow2.path().to_str().context("Invalid UTF-8")?;
    cmd(&[
        "qemu-img", "convert", "-f", "raw", "-O", "qcow2", raw_path, qcow2_path,
    ])?;

    Ok(qcow2)
}
fn build_qcow2(image: &str, mount_path: &Path) -> Result<()> {
    let reg_dir = registry_dir(image)?;
    fs::create_dir_all(&reg_dir)?;
    println!("Registry: {}", reg_dir.display());

    let qcow2 = to_qcow2_ext4(mount_path, &reg_dir)?;

    fs::copy(qcow2.path(), reg_dir.join("image.qcow2"))
        .context("Failed to copy qcow2 image to registry")?;

    Ok(())
}

fn build_initrd(image: &str, mount_path: &Path) -> Result<()> {
    let mount_path = mount_path.to_str().context("mount path is not UTF-8")?;

    let reg_dir = registry_dir(image)?;
    fs::create_dir_all(&reg_dir)?;
    println!("Registry: {}", reg_dir.display());

    // Pack the rootfs as a gzipped newc cpio archive. The cpio pipeline
    // runs inside the namespace, its stdout is compressed here.
    let initrd_path = reg_dir.join("initrd");
    let script = r#"cd "$1" && find . | cpio -o -H newc"#;
    let mut cpio = Command::new("podman")
        .args(["unshare", "--", "sh", "-c", script, "sh", mount_path])
        .stdout(Stdio::piped())
        .spawn()
        .context("Failed to spawn the cpio pipeline")?;

    let initrd_file = File::create(&initrd_path).context("Failed to create initrd file")?;
    let mut encoder = GzEncoder::new(initrd_file, flate2::Compression::default());
    let cpio_stdout = cpio.stdout.take().context("cpio stdout was not piped")?;
    let copied = io::copy(&mut io::BufReader::new(cpio_stdout), &mut encoder);
    let stream_result = copied
        .and_then(move |_| encoder.finish().map(|_| ()))
        .context("Failed to write the gzipped initrd");
    let cpio_status = cpio
        .wait()
        .context("Failed to wait for the cpio pipeline")?;

    if let Err(e) = stream_result {
        let _ = fs::remove_file(&initrd_path);
        return Err(e);
    }
    if !cpio_status.success() {
        let _ = fs::remove_file(&initrd_path);
        bail!("find/cpio pipeline failed with status {cpio_status}");
    }
    println!("Created initrd -> {}", initrd_path.display());

    Ok(())
}