proofman-starks-lib-c 1.1.0-alpha

use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;

/// Detects whether GPU (CUDA) support is available.
/// Returns false if the cpu-only feature is set or if no CUDA toolkit is found.
fn detect_gpu() -> bool {
    if cfg!(feature = "cpu-only") {
        return false;
    }
    // Check for nvcc in standard CUDA location or PATH
    let nvcc_in_cuda = Path::new("/usr/local/cuda/bin/nvcc").exists();
    let nvcc_in_path = Command::new("nvcc").arg("--version").output().map(|o| o.status.success()).unwrap_or(false);
    nvcc_in_cuda || nvcc_in_path
}

fn main() {
    // CUDA arch resolution lives entirely in pil2-stark/Makefile
    println!("cargo:rerun-if-env-changed=CUDA_ARCHS");
    println!("cargo:rerun-if-env-changed=CUDA_ARCH");
    println!("cargo:rerun-if-env-changed=CUDA_GENCODE_FLAGS");

    // Determine if GPU support should be used:
    // - If cpu-only feature is set, always use CPU
    // - Otherwise, auto-detect CUDA availability
    let use_gpu = if cfg!(feature = "cpu-only") {
        println!("cargo:warning=[BUILD INFO] STARKS compiled with CPU-only support (feature enabled)");
        false
    } else if detect_gpu() {
        println!("cargo:warning=[BUILD INFO] STARKS compiled with GPU support");
        true
    } else {
        println!("cargo:warning=[BUILD INFO] STARKS compiled with CPU-only support (CUDA not detected)");
        false
    };

    // Set build mode as environment variable for runtime access
    if use_gpu {
        println!("cargo:rustc-env=STARKS_BUILD_MODE=GPU");
    } else {
        println!("cargo:rustc-env=STARKS_BUILD_MODE=CPU");
    }

    // Sources are carried by the `proofman-starks-src` crate. Build in place when
    // that tree is writable (a local path dependency / this workspace); when it is
    // a read-only registry checkout (a published consumer), mirror it into OUT_DIR
    // first — the Makefile generates files in-tree (goldilocks `configure.sh`,
    // `.simd_stamp`) and writes `build/` and `lib/` alongside the sources.
    let vendored = proofman_starks_src::source_dir();
    // Writable tree = local path dependency (this workspace); read-only = a
    // published registry checkout that must be mirrored before building.
    let in_place = is_writable(&vendored);
    let pil2_stark_path = if in_place {
        vendored.canonicalize().unwrap_or(vendored)
    } else {
        let dst = Path::new(&env::var("OUT_DIR").unwrap()).join("pil2-stark");
        sync_tree(&vendored, &dst);
        dst.canonicalize().unwrap_or(dst)
    };
    let library_folder = if use_gpu { pil2_stark_path.join("lib-gpu") } else { pil2_stark_path.join("lib") };
    let library_name = if use_gpu { "starksgpu" } else { "starks" };
    let lib_file = library_folder.join(format!("lib{library_name}.a"));

    // The CPU and GPU variants of this script can run concurrently against the
    // same checkout (one cargo build graph may contain both, and a second cargo
    // process or rust-analyzer can overlap with either). All staleness probes,
    // cleans and make invocations below mutate shared in-tree state, so they
    // must be serialized; the lock is held until this process exits.
    let lock_path = pil2_stark_path.join(".build_lock");
    let build_lock = fs::File::options()
        .create(true)
        .truncate(false)
        .write(true)
        .open(&lock_path)
        .unwrap_or_else(|e| panic!("Failed to open build lock {}: {e}", lock_path.display()));
    build_lock.lock().unwrap_or_else(|e| panic!("Failed to acquire build lock {}: {e}", lock_path.display()));

    // For GPU builds: in a local checkout, auto-initialize the blst/sppark
    // submodules exactly as before (a fresh clone may not have them yet). In a
    // published checkout their files are bundled into the crate, so skip the git
    // call — there is no repo there. Then compile blst.
    if use_gpu {
        if in_place {
            ensure_gpu_submodules_initialized(&pil2_stark_path);
        }
        ensure_blst_compiled(&pil2_stark_path);
    }

    let tracked_files = find_tracked_files(&pil2_stark_path);
    for file in &tracked_files {
        println!("cargo:rerun-if-changed={}", file.display());
    }
    println!("cargo:rerun-if-changed={}", lib_file.display());

    // Detect if the Makefile changed since the last build (see the clean below).
    let makefile_path = pil2_stark_path.join("Makefile");
    let makefile_stamp_path = library_folder.join(".makefile_stamp");
    let current_makefile = fs::read(&makefile_path).ok();
    let stored_makefile = fs::read(&makefile_stamp_path).ok();
    let makefile_changed = current_makefile.is_some() && current_makefile != stored_makefile;

    // No staleness gate here: cargo only re-runs this script when one of the
    // `rerun-if-changed` / `rerun-if-env-changed` inputs above actually moved
    let target = if use_gpu { "starks_lib_gpu" } else { "starks_lib" };

    // Clean build when the Makefile itself changes: compiler flag edits (e.g.
    // toggling -D__AVX512__) aren't tracked by make's .d files, so a flag flip
    // would otherwise leave stale objects linked into the new library.
    if makefile_changed {
        eprintln!("Makefile changed — running clean rebuild...");
        // Variant-scoped: a full `make clean` would delete the other variant's
        // build dirs, which may belong to a build that just finished or
        // (without the lock) one still in flight.
        run_command("make", &[if use_gpu { "clean_gpu" } else { "clean_cpu" }], &pil2_stark_path);
    }
    eprintln!("Running make -j {target}...");
    run_command("make", &["-j", target], &pil2_stark_path);

    // Write the stamp after make succeeds (make creates the output directory).
    if let Some(content) = &current_makefile {
        if let Err(e) = fs::write(&makefile_stamp_path, content) {
            eprintln!(
                "Warning: failed to write Makefile stamp {:?}: {e} — next build will recompile",
                makefile_stamp_path
            );
        }
    }

    // Absolute path to the library
    let abs_lib_path = library_folder.canonicalize().unwrap_or_else(|_| library_folder.clone());

    if !lib_file.exists() {
        if use_gpu {
            panic!("`libstarksgpu.a` was not found at {}", lib_file.display());
        } else {
            panic!("`libstarks.a` was not found at {}", lib_file.display());
        }
    }

    // Add platform-specific library search paths
    if cfg!(target_os = "macos") {
        // Get Homebrew prefix for macOS
        let homebrew_prefix = Command::new("brew")
            .arg("--prefix")
            .output()
            .map(|output| String::from_utf8_lossy(&output.stdout).trim().to_string())
            .unwrap_or_else(|_| "/opt/homebrew".to_string()); // Default for Apple Silicon

        println!("cargo:rustc-link-search=native={homebrew_prefix}/lib");
        println!("cargo:rustc-link-search=native={homebrew_prefix}/opt/libomp/lib");
        println!("cargo:rustc-link-search=native={homebrew_prefix}/opt/libsodium/lib");
        println!("cargo:rustc-link-search=native={homebrew_prefix}/opt/gmp/lib");
        println!("cargo:rustc-link-search=native={homebrew_prefix}/opt/openssl/lib");

        // Also add system paths
        println!("cargo:rustc-link-search=native=/Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk/usr/lib");
    } else if cfg!(target_os = "linux") {
        // Standard Linux library paths
        println!("cargo:rustc-link-search=native=/usr/lib");
        println!("cargo:rustc-link-search=native=/usr/local/lib");
        println!("cargo:rustc-link-search=native=/usr/lib/x86_64-linux-gnu");
    }

    // Link the static library
    println!("cargo:rustc-link-search=native={}", abs_lib_path.display());
    println!("cargo:rustc-link-lib=static={library_name}");
    if use_gpu {
        // Add the CUDA library path
        let cuda_path = "/usr/local/cuda/lib64"; // Adjust this path if necessary
        println!("cargo:rustc-link-search=native={cuda_path}");
        println!("cargo:rustc-link-lib=static=cudart_static"); // Link the CUDA runtime library statically
                                                               // cudart_static requires additional system libraries
        println!("cargo:rustc-link-lib=dylib=dl");
        println!("cargo:rustc-link-lib=dylib=rt");

        // Add the blst library for GPU MSM
        let blst_path = pil2_stark_path.join("external/blst");
        let blst_lib_path = blst_path.canonicalize().unwrap_or_else(|_| blst_path.clone());
        println!("cargo:rustc-link-search=native={}", blst_lib_path.display());
        println!("cargo:rustc-link-lib=static=blst");
    }

    // Link required libraries with platform-specific handling
    if cfg!(target_os = "macos") {
        // macOS library linking (matches Makefile LDFLAGS)
        for lib in &["sodium", "pthread", "gmp", "gmpxx", "c++", "omp"] {
            println!("cargo:rustc-link-lib={lib}");
        }
    } else {
        // Linux library linking
        for lib in &["sodium", "pthread", "gmp", "stdc++", "gmpxx", "crypto", "iomp5"] {
            println!("cargo:rustc-link-lib={lib}");
        }
        // libstarks.a is always compiled with -D__USE_MPI_RMA__ on Linux, so link MPI
        println!("cargo:rustc-link-search=native=/usr/lib/x86_64-linux-gnu/openmpi/lib");
        println!("cargo:rustc-link-lib=mpi");
    }
}

/// Runs an external command and checks for errors
fn run_command(cmd: &str, args: &[&str], dir: &Path) {
    let status = Command::new(cmd)
        .args(args)
        .current_dir(dir)
        .status()
        .unwrap_or_else(|e| panic!("Failed to execute `{cmd}`: {e}"));

    if !status.success() {
        panic!("Command `{}` failed with exit code {:?}", cmd, status.code());
    }
}

/// Recursively finds all files in `pil2-stark`, skipping build output directories.
fn find_tracked_files(dir: &Path) -> Vec<PathBuf> {
    let mut files = Vec::new();
    if let Some(name) = dir.file_name().and_then(|n| n.to_str()) {
        if matches!(name, "build" | "build-gpu" | "build_gpu" | "lib" | "lib-gpu" | ".vscode" | ".git") {
            return files;
        }
    }
    if let Ok(entries) = fs::read_dir(dir) {
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                files.extend(find_tracked_files(&path));
            } else {
                // Skip build-generated files: .mk (make includes), .d (dependency
                // files), and the build lock
                let ext = path.extension().and_then(|e| e.to_str());
                if matches!(ext, Some("mk" | "d")) {
                    continue;
                }
                let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
                if name == ".build_lock" || name.ends_with("_stamp") {
                    continue;
                }
                files.push(path);
            }
        }
    }
    files
}

/// Ensures GPU-required submodules (blst and sppark) are initialized. Local
/// checkouts only — a published crate ships the submodule files in-tree, so the
/// caller skips this (there is no git repo to update there).
fn ensure_gpu_submodules_initialized(pil2_stark_path: &Path) {
    let blst_path = pil2_stark_path.join("external/blst");
    let sppark_path = pil2_stark_path.join("external/sppark");

    if !is_submodule_initialized(&blst_path) || !is_submodule_initialized(&sppark_path) {
        eprintln!("GPU submodules not fully initialized, running git submodule update...");
        let workspace_root = pil2_stark_path.parent().unwrap_or(pil2_stark_path);
        run_command("git", &["submodule", "update", "--init", "--recursive"], workspace_root);
    }
}

/// Whether the source tree is writable — true for a local path dependency (this
/// workspace), false for a read-only registry checkout of the vendored sources.
///
/// Probes by opening an existing file (the Makefile) for writing rather than
/// creating a temp file, so it leaves no trace in the source tree. Opening for
/// write without writing changes neither the file's contents nor its mtime.
fn is_writable(dir: &Path) -> bool {
    fs::OpenOptions::new().write(true).open(dir.join("Makefile")).is_ok()
}

/// Grant the owner write permission without making the file world-writable
/// (`set_readonly(false)` would set every write bit on Unix).
#[cfg(unix)]
fn make_writable(mut perms: fs::Permissions) -> fs::Permissions {
    use std::os::unix::fs::PermissionsExt;
    perms.set_mode(perms.mode() | 0o200);
    perms
}

#[cfg(not(unix))]
fn make_writable(mut perms: fs::Permissions) -> fs::Permissions {
    #[allow(clippy::permissions_set_readonly_false)]
    perms.set_readonly(false);
    perms
}

/// Recursively mirror `src` into `dst`, copying a file only when the source is
/// newer (preserving mtime) so the Makefile's incremental rebuilds keep working
/// across cargo invocations. Used only for read-only (published) source trees.
fn sync_tree(src: &Path, dst: &Path) {
    fs::create_dir_all(dst).expect("create OUT_DIR source mirror");
    for entry in fs::read_dir(src).expect("read vendored source dir").flatten() {
        let from = entry.path();
        let to = dst.join(entry.file_name());
        if from.is_dir() {
            sync_tree(&from, &to);
        } else if from.is_file() {
            let src_mtime = fs::metadata(&from).and_then(|m| m.modified()).ok();
            let dst_mtime = fs::metadata(&to).and_then(|m| m.modified()).ok();
            let stale = match (src_mtime, dst_mtime) {
                (Some(s), Some(d)) => s > d,
                _ => true,
            };
            if stale {
                fs::copy(&from, &to).expect("copy vendored source file");
                // fs::copy preserves source permissions; a published source is a
                // read-only registry checkout, but the build must rewrite this
                // mirror (configure.sh codegen, make objects, stamps), so add the
                // owner write bit to the copy.
                if let Ok(meta) = fs::metadata(&to) {
                    let perms = meta.permissions();
                    if perms.readonly() {
                        let _ = fs::set_permissions(&to, make_writable(perms));
                    }
                }
                if let (Some(s), Ok(f)) = (src_mtime, fs::File::options().write(true).open(&to)) {
                    let _ = f.set_modified(s);
                }
            }
        }
    }
}

/// Ensures the blst library is compiled for GPU builds
fn ensure_blst_compiled(pil2_stark_path: &Path) {
    let blst_path = pil2_stark_path.join("external/blst");
    let blst_lib = blst_path.join("libblst.a");

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

    if blst_lib.exists() {
        eprintln!("blst library already exists at {}", blst_lib.display());
        return;
    }

    eprintln!("blst library not found at {}, compiling...", blst_lib.display());

    let build_script = blst_path.join("build.sh");

    // Track blst build script and source files for changes
    println!("cargo:rerun-if-changed={}", build_script.display());
    println!("cargo:rerun-if-changed={}", blst_path.join("src").display());
    println!("cargo:rerun-if-changed={}", blst_path.join("build").display());
    if !build_script.exists() {
        panic!("blst build.sh not found at {} — vendored blst sources are incomplete.", build_script.display());
    }

    // Run the blst build script
    let status = Command::new("sh")
        .arg("build.sh")
        .current_dir(&blst_path)
        .status()
        .unwrap_or_else(|e| panic!("Failed to execute blst build.sh: {e}"));

    if !status.success() {
        panic!("blst build.sh failed with exit code {:?}", status.code());
    }

    // Verify the library was created
    if !blst_lib.exists() {
        panic!("blst compilation completed but libblst.a was not created at {}", blst_lib.display());
    }

    eprintln!("blst library successfully compiled at {}", blst_lib.display());
}

/// Checks if a git submodule is initialized (has a .git entry or non-empty dir).
fn is_submodule_initialized(path: &Path) -> bool {
    // Initialized submodules have a .git file (not directory) pointing to parent's .git/modules/
    if path.join(".git").exists() {
        return true;
    }
    // Fallback: check if directory exists and is not empty
    if let Ok(mut entries) = fs::read_dir(path) {
        return entries.next().is_some();
    }
    false
}