use std::process::Command;

fn main() {
    println!("cargo:rerun-if-changed=kernels/");

    if std::env::var("CARGO_FEATURE_CUDA").is_ok() {
        compile_kernels();
    }
}

fn compile_kernels() {
    let out_dir = std::env::var("OUT_DIR").unwrap();

    let kernels = ["test_kernels"];

    for name in kernels {
        let src = format!("kernels/{}.cu", name);
        let dst = format!("{}/{}.ptx", out_dir, name);

        println!("cargo:rerun-if-changed={}", src);

        let status = Command::new("nvcc")
            .args([
                "-ptx",
                "-o",
                &dst,
                &src,
                "--gpu-architecture=sm_75", // Adjust for your GPU
                "-O3",
                "--use_fast_math",
            ])
            .status()
            .expect("Failed to run nvcc");

        if !status.success() {
            panic!("nvcc failed on {}", src);
        }

        println!("cargo:warning=Compiled {} -> {}", src, dst);
    }

    println!("cargo:rustc-env=KERNEL_PTX_DIR={}", out_dir);
}