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

#[derive(Parser, Debug)]
#[command(name = "slim")]
#[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(())
}

fn registry_dir(image: &str) -> PathBuf {
    let base = xdg::BaseDirectories::with_prefix("slim-rs").unwrap();
    base.place_data_file(format!("registry/{}/vmlinuz", image)).unwrap();
    base.place_data_file(format!("registry/{}/initrd", image)).unwrap().parent().unwrap().to_path_buf()
}

fn build(image: &str) -> Result<()> {
    // Step 1: mount via podman unshare (verified in Slice 0)
    let output = Command::new("podman")
        .args(["unshare", "--", "podman", "image", "mount", image])
        .output()
        .context("Failed to run podman unshare image mount")?;
    let mount_path_str = String::from_utf8_lossy(&output.stdout).trim().to_string();
    if mount_path_str.is_empty() {
        anyhow::bail!("podman image mount returned empty path");
    }
    let mount_path = Path::new(&mount_path_str);
    println!("Mounted at: {}", mount_path.display());

    // Step 2: find vmlinuz
    let vmlinuz_src = mount_path.join("vmlinuz");
    if !vmlinuz_src.exists() {
        anyhow::bail!("Image missing required /vmlinuz. Place kernel at /vmlinuz in Containerfile.");
    }

    // Step 3: prepare registry dir
    let reg_dir = registry_dir(image);
    fs::create_dir_all(&reg_dir)?;
    println!("Registry: {}", reg_dir.display());

    // Copy vmlinuz out
    let vmlinuz_dst = reg_dir.join("vmlinuz");
    fs::copy(&vmlinuz_src, &vmlinuz_dst).context("Copy vmlinuz to registry")?;
    println!("Copied vmlinuz -> {}", vmlinuz_dst.display());

    // Step 4: create initrd via cpio binary (using system cpio for reliability)
    let initrd_raw = reg_dir.join("initrd.tmp");
    let initrd_path = reg_dir.join("initrd");

    // Build cpio from rootfs (excluding vmlinuz since it is separate kernel)
    let mut find_cmd = Command::new("find")
        .arg(".")
        .current_dir(mount_path)
        .stdout(Stdio::piped())
        .spawn()
        .context("Failed to spawn find")?;
    let find_stdout = find_cmd.stdout.take().unwrap();

    let mut cpio_cmd = Command::new("cpio")
        .args(["-o", "-H", "newc"])
        .current_dir(mount_path)
        .stdin(Stdio::from(find_stdout))
        .stdout(Stdio::from(File::create(&initrd_raw)?))
        .spawn()
        .context("Failed to spawn cpio")?;

    let find_status = find_cmd.wait()?;
    let cpio_status = cpio_cmd.wait()?;
    if !find_status.success() || !cpio_status.success() {
        anyhow::bail!("find or cpio failed");
    }

    // Gzip initrd
    let initrd_file = fs::File::open(&initrd_raw)?;
    let mut initrd_writer = flate2::write::GzEncoder::new(File::create(&initrd_path)?, flate2::Compression::default());
    std::io::copy(&mut std::io::BufReader::new(initrd_file), &mut initrd_writer)?;
    initrd_writer.finish()?;
    fs::remove_file(&initrd_raw)?;
    println!("Created initrd -> {}", initrd_path.display());

    // Unmount image
    let _ = Command::new("podman")
        .args(["unshare", "--", "podman", "image", "unmount", image])
        .output();
    println!("Unmounted image.");

    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() {
        anyhow::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()
    );
    // Launch with a 5-second timeout for quick smoke verification
    let status = Command::new("timeout")
        .args(["5", "qemu-system-x86_64",
               "-m", "256",
               "-nographic",
               "-kernel", vmlinuz_path.to_str().unwrap(),
               "-initrd", initrd_path.to_str().unwrap(),
               "-append", "console=ttyS0"])
        .status()
        .context("Failed to launch qemu-system-x86_64")?;
    if status.success() {
        println!("QEMU exited cleanly (timeout or normal exit).");
    } else {
        println!("QEMU exited with non-zero status (timeout/expected for initrd boot without init).");
    }
    Ok(())
}