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

#[derive(Parser, Debug)]
#[command(about = "Build bootable initrd VMs from container images")]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand, Debug)]
enum Commands {
    /// Build initrd from an image (podman image must exist)
    Build {
        /// Image name (e.g. alpine:3.15 or slim:tag)
        image: String,
        /// Size hint (unused in v1 initrd)
        #[arg(short = 's', long, default_value = "512")]
        _size: String,
    },
    /// Launch the VM via QEMU using registry artifacts
    Run {
        /// Image tag / registry subdir name
        name: String,
    },
}

fn main() -> Result<()> {
    let cli = Cli::parse();
    match cli.command {
        Commands::Build { image, _size: _ } => {
            build(&image)?;
        }
        Commands::Run { name } => {
            run(&name)?;
        }
    }
    Ok(())
}

/// Validate a user-supplied image name before it is used in a registry path.
///
/// The name becomes a directory under `$XDG_DATA_HOME/slim-rs/registry/`, so
/// path separators and `.`/`..` would allow escaping that directory.
/// Allow-list: ASCII alphanumerics plus `.`, `_`, `:`, `-` (covers image
/// refs like `alpine:3.15`).
fn validate_image_name(name: &str) -> Result<()> {
    let valid = !name.is_empty()
        && name != "."
        && name != ".."
        && name
            .chars()
            .all(|c| c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | ':' | '-'));
    if !valid {
        bail!(
            "invalid image name '{name}': must be non-empty and contain only \
             alphanumerics, '.', '_', ':' and '-' (no path separators or '..')"
        );
    }
    Ok(())
}

fn registry_dir(image: &str) -> Result<PathBuf> {
    validate_image_name(image)?;
    let base = xdg::BaseDirectories::with_prefix("slim-rs");
    let initrd_path = base
        .place_data_file(format!("registry/{image}/initrd"))
        .context("Failed to write to XDG_DATA_HOME")?;
    let reg_dir = initrd_path
        .parent()
        .expect("data file has a parent")
        .to_owned();

    // Defense in depth: the constructed path must stay inside the registry.
    let registry_root = base
        .place_data_file("registry/.root-check")
        .context("Failed to write to XDG_DATA_HOME")?
        .parent()
        .expect("data file has a parent")
        .to_owned();
    if !reg_dir.starts_with(&registry_root) {
        bail!(
            "invalid image name '{image}': registry path escapes {}",
            registry_root.display()
        );
    }
    Ok(reg_dir)
}

fn build(image: &str) -> Result<()> {
    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 = build_artifacts(image, &mount_path);

    let unmounted = Command::new("podman")
        .args(["unshare", "--", "podman", "image", "unmount", image])
        .output()
        .is_ok_and(|out| out.status.success());
    if unmounted {
        println!("Unmounted image.");
    } else {
        eprintln!("warning: failed to unmount image '{}'", image);
    }

    result
}

fn mount_image(image: &str) -> Result<PathBuf> {
    let output = Command::new("podman")
        .args(["unshare", "--", "podman", "image", "mount", image])
        .output()
        .context("Failed to run podman image mount")?;
    let mount_path = String::from_utf8_lossy(&output.stdout).trim().to_string();
    if !output.status.success() || mount_path.is_empty() {
        bail!(
            "podman image mount failed: {}",
            String::from_utf8_lossy(&output.stderr).trim()
        );
    }
    Ok(PathBuf::from(mount_path))
}

fn build_artifacts(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 = Command::new("podman")
        .args(["unshare", "--", "test", "-f", &vmlinuz_src])
        .status()
        .context("Failed to check for /vmlinuz inside the image mount")?;
    if !has_kernel.success() {
        bail!("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");
    let cp_status = Command::new("podman")
        .args([
            "unshare",
            "--",
            "cp",
            &vmlinuz_src,
            vmlinuz_dst.to_str().context("registry path is not UTF-8")?,
        ])
        .status()
        .context("Failed to copy vmlinuz out of the image mount")?;
    if !cp_status.success() {
        bail!("Failed to copy vmlinuz to registry");
    }
    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(())
}

fn run(name: &str) -> Result<()> {
    let reg_dir = registry_dir(name)?;
    let vmlinuz_path = reg_dir.join("vmlinuz");
    let initrd_path = reg_dir.join("initrd");
    if !vmlinuz_path.exists() || !initrd_path.exists() {
        bail!(
            "Registry missing vmlinuz/initrd for '{}'. Run `slim build` first.",
            name
        );
    }
    println!("Booting {} from registry: {}", name, reg_dir.display());
    println!(
        "  qemu-system-x86_64 -m 256 -nographic -kernel {} -initrd {} -append 'console=ttyS0'",
        vmlinuz_path.display(),
        initrd_path.display()
    );

    let cmdline = [
        "console=ttyS0,115200",
        "root=/dev/ram0",
        "rw",
        "earlyprintk=serial",
        "nokaslr",
        "raid=noautodetect",
    ]
    .join(" ");

    // Launch with a 5-second timeout for quick smoke verification
    let status = Command::new("qemu-system-x86_64")
        .args(["-m", "1024M"])
        .arg("-kernel")
        .arg(vmlinuz_path.as_os_str())
        .arg("-initrd")
        .arg(initrd_path.as_os_str())
        .args(["-append", &cmdline])
        .args(["-nographic"])
        .status()
        .context("Failed to launch qemu-system-x86_64")?;

    match status.code() {
        Some(0) => println!("QEMU exited cleanly."),
        Some(126) | Some(127) => {
            bail!("Failed to start qemu-system-x86_64 (is it installed and in PATH?)")
        }
        Some(code) => bail!("QEMU exited with code {code}"),
        None => bail!("QEMU terminated by a signal"),
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::validate_image_name;

    #[test]
    fn accepts_valid_image_names() {
        for name in [
            "alpine:3.15",
            "slim",
            "mock-vm",
            "my_vm",
            "a.b.c",
            "vm-1.2.3",
        ] {
            validate_image_name(name).unwrap_or_else(|e| panic!("{name} rejected: {e}"));
        }
    }

    #[test]
    fn rejects_traversal_and_separators() {
        for name in [
            "",
            ".",
            "..",
            "../evil",
            "../../tmp/evil",
            "a/b",
            "/abs",
            "a\\b",
            "foo bar",
            "vm;rm",
            "vm\n",
        ] {
            assert!(
                validate_image_name(name).is_err(),
                "{name:?} should be rejected"
            );
        }
    }
}