sp1-gpu-sys 6.0.0-beta.1

use std::path::PathBuf;
use std::{env, fs};

fn cbindgen_builder() -> cbindgen::Builder {
    /// The warning placed in the cbindgen header.
    const AUTOGEN_WARNING: &str =
        "/* Automatically generated by `cbindgen`. Not intended for manual editing. */";

    // Generate a header containing bindings to the crate.
    cbindgen::Builder::new()
        .with_pragma_once(true)
        .with_autogen_warning(AUTOGEN_WARNING)
        .with_no_includes()
        .with_sys_include("cstdint")
        .with_parse_deps(true)
        .with_parse_include(&[
            "sp1-hypercube",
            "sp1-primitives",
            "p3-koala-bear",
            "sp1-core-machine",
            "slop-koala-bear",
            "sp1-core-executor",
            "sp1-recursion-executor",
            "sp1-recursion-machine",
        ])
        .with_parse_extra_bindings(&[
            "sp1-hypercube",
            "p3-koala-bear",
            "slop-koala-bear",
            "sp1-recursion-executor",
            "sp1-recursion-machine",
            "sp1-core-machine",
        ])
        .rename_item("KoalaBear", "KoalaBearP3")
        .include_item("GlobalInteractionEvent")
        .include_item("GlobalCols")
        .include_item("MemoryAccessColsChips")
        .include_item("Poseidon2Instr")
        .include_item("Poseidon2PreprocessedColsWide")
        .include_item("Poseidon2Event")
        .include_item("SelectInstr")
        .include_item("SelectPreprocessedCols")
        .include_item("SelectEvent")
        .include_item("SelectCols")
        .include_item("BaseAluInstr")
        .include_item("BaseAluAccessCols")
        .include_item("BaseAluEvent")
        .include_item("BaseAluValueCols")
        .include_item("ExtAluInstr")
        .include_item("ExtAluAccessCols")
        .include_item("ExtAluEvent")
        .include_item("ExtAluValueCols")
        .include_item("ExtFeltInstr")
        .include_item("ExtFeltEvent")
        .include_item("ConvertAccessCols")
        .include_item("ConvertValueCols")
        .include_item("Poseidon2LinearLayerAccessCols")
        .include_item("Poseidon2LinearLayerValueCols")
        .include_item("Poseidon2LinearLayerInstr")
        .include_item("Poseidon2SBoxInstr")
        .include_item("Poseidon2SBoxIo")
        .include_item("Poseidon2SBoxAccessCols")
        .include_item("Poseidon2SBoxValueCols")
        .include_item("NUM_CONVERT_ENTRIES_PER_ROW")
        .include_item("NUM_LINEAR_ENTRIES_PER_ROW")
        .include_item("NUM_SBOX_ENTRIES_PER_ROW")
        .include_item("PERMUTATION_WIDTH")
        .include_item("PrefixSumChecksEvent")
        .include_item("PrefixSumChecksCols")
        .with_namespace("sp1_gpu_sys")
        .with_crate(env::var("CARGO_MANIFEST_DIR").unwrap())
}

/// Place a relative symlink pointing to `original` at `link`.
fn rel_symlink_file<P, Q>(original: P, link: Q)
where
    P: AsRef<std::path::Path>,
    Q: AsRef<std::path::Path>,
{
    #[cfg(unix)]
    use std::os::unix::fs::symlink;
    #[cfg(windows)]
    use std::os::windows::fs::symlink_file as symlink;

    let target_dir = link.as_ref().parent().unwrap();
    fs::create_dir_all(target_dir).unwrap();
    let _ = fs::remove_file(&link);
    let relpath = pathdiff::diff_paths(original, target_dir).unwrap();
    symlink(relpath, link).unwrap();
}

/// Check if CUDA is available on this system.
fn detect_cuda() -> bool {
    // Track rerun-if-env-changed for the explicit override
    println!("cargo:rerun-if-env-changed=SP1_CUDA_ENABLED");

    // 1. Check explicit SP1_CUDA_ENABLED env var
    if let Ok(val) = env::var("SP1_CUDA_ENABLED") {
        match val.to_lowercase().as_str() {
            "0" | "false" => return false,
            "1" | "true" => return true,
            _ => {} // Fall through to auto-detection
        }
    }

    // 2. Try running nvcc --version
    if std::process::Command::new("nvcc")
        .arg("--version")
        .output()
        .map(|o| o.status.success())
        .unwrap_or(false)
    {
        return true;
    }

    // 3. Check CUDA_PATH env var
    if env::var("CUDA_PATH").is_ok() {
        return true;
    }

    // 4. Check common CUDA installation paths
    if std::path::Path::new("/usr/local/cuda/bin/nvcc").exists() {
        return true;
    }
    if std::path::Path::new("/opt/cuda/bin/nvcc").exists() {
        return true;
    }

    false
}

fn main() {
    // Directives for tracking changes in folders
    println!("cargo:rerun-if-changed=../../include/");
    println!("cargo:rerun-if-changed=../../lib/");
    println!("cargo:rerun-if-changed=../../sppark/");
    println!("cargo:rerun-if-changed=../../CMakeLists.txt");
    println!("cargo:rerun-if-changed=CMakeLists.txt");
    // Directives for tracking changes in compilation profiles
    println!("cargo:rerun-if-env-changed=PROFILE");
    println!("cargo:rerun-if-env-changed=OPT_LEVEL");
    println!("cargo:rerun-if-env-changed=DEBUG");
    println!("cargo:rerun-if-env-changed=CUDA_ARCHS");
    println!("cargo:rerun-if-env-changed=PROFILE_DEBUG_DATA");

    // The crate directory.
    let crate_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());

    // The output directory, where built artifacts should be placed.
    let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());

    // The directory to place headers into.
    let out_include_dir = out_dir.join("include");
    const INCLUDE_DIRNAME: &str = "include";

    // Generate cbindgen headers BEFORE building with CMake
    let cbindgen_hpp = "sp1-gpu-cbindgen.hpp";
    let header_path = out_include_dir.join(cbindgen_hpp);

    match cbindgen_builder().generate() {
        Ok(bindings) => {
            fs::create_dir_all(&out_include_dir).unwrap();
            if bindings.write_to_file(&header_path) {
                // Symlink header into `target/include` for IDE purposes
                let mut dir = out_dir.clone();
                while !dir.ends_with("target") {
                    if !dir.pop() {
                        break;
                    }
                }
                if dir.ends_with("target") {
                    rel_symlink_file(&header_path, dir.join(INCLUDE_DIRNAME).join(cbindgen_hpp));
                }
            }
        }
        Err(cbindgen::Error::ParseSyntaxError { .. }) => {
            // Ignore parse errors in the build script.
            // We let rust-analyzer/rustc run to report the syntax error with diagnostics.
            return;
        }
        Err(e) => panic!("{e:?}"),
    }

    // Check if CUDA is available before attempting to build
    if !detect_cuda() {
        println!("cargo:warning=CUDA not detected, skipping GPU build");
        return;
    }

    // Build using CMake
    let mut cmake_config = cmake::Config::new(".");

    // Export compile commands for clangd IDE support
    cmake_config.define("CMAKE_EXPORT_COMPILE_COMMANDS", "ON");

    // Pass CUDA architectures to CMake only if explicitly set
    // Otherwise, CMake will use its own version-based defaults
    if let Ok(cuda_archs) = env::var("CUDA_ARCHS") {
        cmake_config.define("CUDA_ARCHS", &cuda_archs);
    }

    // Pass cbindgen include directory
    cmake_config.define("CBINDGEN_INCLUDE_DIR", out_include_dir.display().to_string());

    // Pass profile debug flag if set
    if let Some(profile_debug_data) = env::var_os("PROFILE_DEBUG_DATA") {
        if profile_debug_data == "true" {
            cmake_config.define("PROFILE_DEBUG_DATA", "true");
        }
    }

    // Determine build type based on profile
    let profile = env::var("PROFILE").unwrap_or_else(|_| "release".to_string());
    if profile == "debug" {
        cmake_config.profile("Debug");
    } else {
        cmake_config.profile("Release");
    }

    // Build via CMake
    let dst = cmake_config.build();

    // Symlink compile_commands.json to project root for clangd
    let compile_commands_src = dst.join("build").join("compile_commands.json");
    if compile_commands_src.exists() {
        let project_root = crate_dir.join("../..").canonicalize().unwrap();
        rel_symlink_file(&compile_commands_src, project_root.join("compile_commands.json"));
    }

    // Link the library
    println!("cargo:rustc-link-search=native={}/../../target/cuda-build/lib", crate_dir.display());
    println!("cargo:rustc-link-lib=static=sys-cuda");

    // Add CUDA library search paths
    if let Ok(cuda_path) = env::var("CUDA_PATH") {
        println!("cargo:rustc-link-search=native={cuda_path}/lib64");
        println!("cargo:rustc-link-search=native={cuda_path}/lib");
    } else {
        println!("cargo:rustc-link-search=native=/usr/local/cuda/lib64");
    }

    // Link CUDA runtime libraries
    println!("cargo:rustc-link-lib=cudart");
    println!("cargo:rustc-link-lib=cudadevrt");

    // Link system libraries
    println!("cargo:rustc-link-lib=stdc++");
    println!("cargo:rustc-link-lib=gomp");
    println!("cargo:rustc-link-lib=dl");

    // Add include directories for compilation if needed
    println!("cargo:include={}", out_include_dir.display());
}