//! Build: mount a podman image, extract vmlinuz, and pack the rootfs as a
//! gzipped cpio initrd in the registry.

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

use crate::command::cmd;
use crate::meta::Meta;
use crate::registry::registry_dir;

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

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

    /// The path to the `init` program for the VM.
    // TODO: infer init based on `COMMAND`/`ENTRYPOINT`?
    #[clap(long, default_value = "/sbin/init")]
    init: String,
}

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

pub(crate) fn build(BuildCmd { kind, name, init }: BuildCmd) -> Result<()> {
    let image = &name;
    let mount_path = mount_image(image)?;
    println!("Mounted at: {}", mount_path.display());

    // Keep the build result so the image is unmounted even when the build fails.
    let result = match kind {
        ImageKind::Initrd => build_initrd(image, &mount_path),
        ImageKind::Qcow2 => build_qcow2(image, &mount_path),
    };

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

    Meta { init }.save(image)?;

    result
}

fn mount_image(image: &str) -> Result<PathBuf> {
    let mount_path = cmd(&["podman", "unshare", "--", "podman", "image", "mount", image])?;
    Ok(PathBuf::from(mount_path.trim()))
}

/// Copy a directory onto a new raw disk image with EXT4.
fn to_raw_ext4(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 mb = 1024 * 1024;
    let size = used
        + (used / 10) // Add 10%
        + 64 * mb; // Add some margin
    let size = size.to_string();

    let raw_file = NamedTempFile::new()?;
    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) -> Result<NamedTempFile> {
    let raw = to_raw_ext4(mount_path)?;
    let qcow2 = NamedTempFile::new()?;
    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)?;

    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 vmlinuz_src = format!("{mount_path}/vmlinuz");

    // The rootless overlay mount only exists inside podman unshare's user
    // namespace, so every read of the rootfs must run there; pipes still
    // cross the namespace boundary.
    let _has_kernel = cmd(&["podman", "unshare", "--", "test", "-f", &vmlinuz_src])
        .context("Image missing required /vmlinuz. Place kernel at /vmlinuz in Containerfile.")?;

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

    let vmlinuz_dst = reg_dir.join("vmlinuz");
    cmd(&[
        "podman",
        "unshare",
        "--",
        "cp",
        &vmlinuz_src,
        vmlinuz_dst.to_str().context("registry path is not UTF-8")?,
    ])
    .context("Failed to copy vmlinuz out of the image mount")?;
    println!("Copied vmlinuz -> {}", vmlinuz_dst.display());

    // Pack the rootfs as a gzipped newc cpio archive; vmlinuz is excluded
    // because QEMU loads the kernel separately via -kernel. 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 . -not -path ./vmlinuz | 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 = std::io::copy(&mut std::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(())
}