//! Build script for cp2k-rs
//!
//! This script clones and builds CP2K from source and links it with the Rust crate.

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

// ── Target-architecture helpers ──────────────────────────────────────────────

/// Returns the target architecture identifier.
/// Set `CP2K_TARGET_ARCH=v4` to compile for x86-64-v4 (AVX-512).
/// Defaults to `"v3"` (x86-64-v3 = AVX2 / FMA / BMI2 / …).
fn cp2k_target_arch() -> String {
    env::var("CP2K_TARGET_ARCH").unwrap_or_else(|_| "v3".to_string())
}

/// Returns the `-march=` flag for the given architecture string.
fn march_flag(arch: &str) -> String {
    match arch {
        "v4" => "-march=x86-64-v4".to_string(),
        _ => "-march=x86-64-v3".to_string(),
    }
}

/// Returns additional FFTW3 `./configure` flags for SIMD support.
fn fftw3_simd_args(arch: &str) -> Vec<&'static str> {
    match arch {
        "v4" => vec!["--enable-avx2", "--enable-avx512"],
        _ => vec!["--enable-avx2"],
    }
}

/// Returns the stable-directory suffix for the given architecture.
/// The v3 build uses bare names (no suffix) to stay backward-compatible
/// with existing caches.
fn arch_dir_suffix(arch: &str) -> &'static str {
    match arch {
        "v4" => "-v4",
        _ => "",
    }
}

fn main() {
    // Only rebuild if the build script changes
    println!("cargo:rerun-if-changed=build.rs");
    println!("cargo:rerun-if-env-changed=CP2K_TARGET_ARCH");
    println!("cargo:rerun-if-changed=libcp2k.h");

    let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
    let profile = env::var("PROFILE").unwrap_or_else(|_| "release".into());

    // Re-run if the worker binary changes so we re-bundle its libs.
    // If the file doesn't exist yet, Cargo treats it as always-changed,
    // ensuring build.rs re-runs after the worker is first built.
    let worker_bin = PathBuf::from(&manifest_dir)
        .join("target")
        .join(&profile)
        .join("cp2k_rs_worker");
    println!("cargo:rerun-if-changed={}", worker_bin.display());

    // Watch for changes in extensions
    if cfg!(feature = "extended") {
        println!("cargo:rerun-if-changed=extensions/fortran/libcp2k_extended.F90");
        println!("cargo:rerun-if-changed=extensions/fortran/libcp2k_mpi.F90");
        println!("cargo:rerun-if-changed=extensions/include/libcp2k_extended.h");
    }

    // Bundle shared libs next to the worker binary so maturin includes them in
    // the wheel. The worker must already exist (built by a prior cargo build
    // --bin cp2k_rs_worker invocation, or by the bundle-wheel.sh script).
    bundle_worker_libs(&manifest_dir, &profile);
    bundle_cp2k_data(&manifest_dir, &profile);

    // Only build CP2K if bindgen feature is enabled
    if env::var("CARGO_FEATURE_BINDGEN").is_err() {
        println!("cargo:warning=Skipping CP2K build - bindgen feature not enabled");
        return;
    }

    let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());

    // Detect MPI installation
    let mpi_base_dir = detect_mpi_installation();
    configure_mpi_environment(&mpi_base_dir);

    // Check if we should build CP2K or use header-only mode
    let should_build_cp2k =
        env::var("CP2K_RS_BUILD_CP2K").unwrap_or_default() == "1" || cfg!(feature = "build-cp2k");

    if should_build_cp2k {
        println!("cargo:warning=Building CP2K from source - this may take 30+ minutes");

        // Use stable paths (under target/<profile>/) so CP2K source and build
        // survive build.rs hash changes between rebuilds.
        // OUT_DIR = .../target/<profile>/build/cp2k-rs-<hash>/out/
        // ancestors: nth(0)=out, nth(1)=cp2k-rs-<hash>, nth(2)=build, nth(3)=<profile>
        let profile_dir = out_dir.ancestors().nth(3).unwrap_or(&out_dir).to_path_buf();
        let arch = cp2k_target_arch();
        let cp2k_stable_dir = profile_dir.join(format!("cp2k-stable{}", arch_dir_suffix(&arch)));
        std::fs::create_dir_all(&cp2k_stable_dir).expect("Failed to create cp2k-stable dir");

        // Clone CP2K repository if it doesn't exist
        let cp2k_src_dir = clone_cp2k_repository(&cp2k_stable_dir);

        // Resolve static OpenBLAS location from the openblas-src build dependency.
        // openblas-src (links = "openblas") emits cargo:LIBRARY=<dir> and
        // cargo:INCLUDE=<dir>, becoming DEP_OPENBLAS_LIBRARY / DEP_OPENBLAS_INCLUDE.
        // openblas-src builds in-tree so DEP_OPENBLAS_LIBRARY is the directory that
        // contains libopenblas.a directly (not a lib/ subdirectory).
        let openblas_lib_dir = PathBuf::from(env::var("DEP_OPENBLAS_LIBRARY").expect(
            "DEP_OPENBLAS_LIBRARY not set – is openblas-src in [dependencies] with 'static' feature?",
        ));
        let openblas_static_lib = openblas_lib_dir.join("libopenblas.a");
        let openblas_include_dir = PathBuf::from(
            env::var("DEP_OPENBLAS_INCLUDE")
                .unwrap_or_else(|_| openblas_lib_dir.to_string_lossy().into_owned()),
        );
        println!(
            "cargo:warning=Using static OpenBLAS: {}",
            openblas_static_lib.display()
        );

        // Build static FFTW3 (with OpenMP support) from source.
        let fftw3_install = build_fftw3(&out_dir);

        // CP2K 2026.1 requires DBCSR as an external pre-built dependency.
        // Clone, build and install DBCSR first, then pass its prefix to CP2K.
        let dbcsr_install_dir = build_dbcsr(&out_dir, &openblas_static_lib);

        // Build ScaLAPACK statically against our static OpenBLAS so that
        // libscalapack.so (and its transitive libflexiblas.so dep) are eliminated
        // from the final binary's dynamic dependencies.
        let scalapack_install = build_scalapack(&out_dir, &openblas_static_lib);

        // Build directory for out-of-source CMake build (also stable)
        let cp2k_build_dir = cp2k_stable_dir.join("build");
        std::fs::create_dir_all(&cp2k_build_dir).expect("Failed to create CP2K build directory");

        // Configure and build CP2K with CMake
        configure_cp2k_cmake(
            &cp2k_src_dir,
            &cp2k_build_dir,
            &dbcsr_install_dir,
            &openblas_static_lib,
            &openblas_include_dir,
            &fftw3_install,
            &scalapack_install,
        );
        build_cp2k_cmake(&cp2k_build_dir);

        // Build Fortran extensions if extended feature is enabled
        if cfg!(feature = "extended") {
            println!("cargo:warning=Building Fortran extensions for extended interface...");
            build_fortran_extensions(&cp2k_src_dir, &cp2k_build_dir, &out_dir);
        }

        // Find the built libcp2k and configure linking
        let libcp2k_path = find_libcp2k(&cp2k_build_dir);
        let lib_dir = libcp2k_path.parent().unwrap();

        // Configure basic linking (BLAS/LAPACK handled by openblas-src; FFTW3 + ScaLAPACK static)
        configure_linking(
            lib_dir,
            &mpi_base_dir,
            &fftw3_install,
            &openblas_static_lib,
            &scalapack_install,
        );

        // Add specific paths for CP2K and DBCSR libraries
        println!("cargo:rustc-link-search=native={}", lib_dir.display());
        println!("cargo:rustc-link-lib=static=cp2k");

        // Add rpath for runtime library location
        println!("cargo:rustc-link-arg=-Wl,-rpath,{}", lib_dir.display());

        // Link DBCSR library (external install, built before CP2K)
        let dbcsr_lib_dir = dbcsr_install_dir.join("lib");
        if dbcsr_lib_dir.exists() {
            println!("cargo:rustc-link-search=native={}", dbcsr_lib_dir.display());
            println!("cargo:rustc-link-lib=static=dbcsr");
        } else {
            panic!(
                "DBCSR install lib dir not found: {}",
                dbcsr_lib_dir.display()
            );
        }

        // Link the extended library if feature is enabled
        if cfg!(feature = "extended") {
            let extended_lib_dir = out_dir.join("extensions_lib");
            if extended_lib_dir.exists() {
                println!(
                    "cargo:rustc-link-search=native={}",
                    extended_lib_dir.display()
                );
                println!("cargo:rustc-link-lib=static=cp2k_extended");
                println!("cargo:warning=Extended interface library linked successfully");
            }
        }
    } else {
        println!("cargo:warning=Skipping CP2K build - using header-only mode");

        // Configure linking (but without CP2K library for now)
        configure_linking_without_cp2k(&mpi_base_dir);

        // Build Fortran extensions against pre-existing CP2K mod files.
        // This compiles the .o files only; linking against libcp2k.a requires
        // the full build-cp2k flow (OpenBLAS, FFTW3, ScaLAPACK, DBCSR…).
        if cfg!(feature = "extended") {
            build_fortran_extensions_from_mods(&out_dir, &profile);
        }
    }

    // Generate Rust bindings to the C header file
    #[cfg(feature = "bindgen")]
    generate_bindings(&manifest_dir, &out_dir);
}

fn detect_mpi_installation() -> PathBuf {
    // Try to find MPI installation
    if let Ok(output) = Command::new("mpicc").arg("--showme:compile").output()
        && output.status.success()
    {
        let compile_flags = String::from_utf8_lossy(&output.stdout);
        for flag in compile_flags.split_whitespace() {
            if let Some(include_path) = flag.strip_prefix("-I")
                && let Some(parent) = Path::new(include_path).parent()
            {
                return parent.to_path_buf();
            }
        }
    }

    // Fallback: common MPI installation paths
    let common_paths = [
        "/usr/lib64/openmpi",
        "/usr/lib/x86_64-linux-gnu/openmpi",
        "/usr/local",
        "/opt/openmpi",
    ];

    for path in &common_paths {
        let path_buf = PathBuf::from(path);
        if path_buf.join("include").join("mpi.h").exists() {
            return path_buf;
        }
    }

    panic!(
        "Could not find MPI installation. Please ensure OpenMPI is installed and mpicc is in PATH."
    );
}

fn configure_mpi_environment(mpi_base_dir: &Path) {
    // Set environment variables for MPI
    println!("cargo:rustc-env=MPI_ROOT={}", mpi_base_dir.display());

    // Add MPI library paths
    if let Some(lib_dir) = find_mpi_lib_dir(mpi_base_dir) {
        println!("cargo:rustc-link-search=native={}", lib_dir.display());
    }

    // Add MPI include paths for compilation
    let include_dir = mpi_base_dir.join("include");
    if include_dir.exists() {
        println!("cargo:include={}", include_dir.display());
    }
}

fn find_mpi_lib_dir(mpi_base_dir: &Path) -> Option<PathBuf> {
    let lib_candidates = ["lib64", "lib"];

    for lib_subdir in &lib_candidates {
        let lib_dir = mpi_base_dir.join(lib_subdir);
        if lib_dir.exists() {
            return Some(lib_dir);
        }
    }
    None
}

/// Clone, build and install DBCSR (required by CP2K 2026.1+).
/// `openblas_static_lib` is the full path to `libopenblas.a`.
/// Returns the install prefix directory.
fn build_dbcsr(out_dir: &Path, openblas_static_lib: &Path) -> PathBuf {
    // Use a stable path (not the hashed OUT_DIR) so DBCSR survives build.rs changes.
    // OUT_DIR is typically .../target/<profile>/build/cp2k-rs-<hash>/out/
    // ancestors: nth(0)=out, nth(1)=cp2k-rs-<hash>, nth(2)=build, nth(3)=<profile>
    // We go up 3 levels to reach the profile directory (e.g. target/release/)
    let profile_dir = out_dir.ancestors().nth(3).unwrap_or(out_dir).to_path_buf();
    let arch = cp2k_target_arch();
    let stable_dir = profile_dir.join(format!("dbcsr-stable{}", arch_dir_suffix(&arch)));
    std::fs::create_dir_all(&stable_dir).expect("Failed to create dbcsr-stable dir");

    let dbcsr_src_dir = stable_dir.join("src");
    let dbcsr_build_dir = stable_dir.join("build");
    let dbcsr_install_dir = stable_dir.join("install");

    if !dbcsr_src_dir.exists() {
        println!("cargo:warning=Cloning DBCSR repository (v2.9.0)...");
        let output = Command::new("git")
            .args([
                "clone",
                "--depth",
                "1",
                "--branch",
                "v2.9.0",
                "https://github.com/cp2k/dbcsr.git",
                dbcsr_src_dir.to_str().unwrap(),
            ])
            .output()
            .expect("Failed to clone DBCSR repository. Make sure git is installed.");
        if !output.status.success() {
            panic!(
                "Failed to clone DBCSR: {}",
                String::from_utf8_lossy(&output.stderr)
            );
        }
        println!("cargo:warning=DBCSR repository cloned successfully");
    } else {
        println!("cargo:warning=DBCSR repository already exists, reusing");
    }

    std::fs::create_dir_all(&dbcsr_build_dir).expect("Failed to create DBCSR build dir");

    let openmpi_pkgconfig = "/usr/lib64/openmpi/lib/pkgconfig";
    let pkg_config_path = match env::var("PKG_CONFIG_PATH") {
        Ok(existing) => format!("{openmpi_pkgconfig}:{existing}"),
        Err(_) => openmpi_pkgconfig.to_string(),
    };

    // Extend PATH to include ~/.local/bin so that 'fypp' (installed via pip --user)
    // is found by CMake's find_program().
    // Also extend PYTHONPATH so that the fypp script (which may have a shebang pointing
    // to a specific Python version) can import the fypp module from the user's site-packages.
    let extended_pythonpath = {
        let home = env::var("HOME").unwrap_or_else(|_| "/root".to_string());
        let base = env::var("PYTHONPATH").unwrap_or_default();
        // Cover common Python versions; harmless if directories don't exist.
        let candidates = [
            format!("{home}/.local/lib/python3.12/site-packages"),
            format!("{home}/.local/lib/python3.11/site-packages"),
            format!("{home}/.local/lib/python3.10/site-packages"),
            format!("{home}/.local/lib/python3.9/site-packages"),
        ];
        let extra = candidates.join(":");
        if base.is_empty() {
            extra
        } else {
            format!("{extra}:{base}")
        }
    };
    let extended_path = {
        let home = env::var("HOME").unwrap_or_else(|_| "/root".to_string());
        let local_bin = format!("{home}/.local/bin");
        // Also check for fypp via `which fypp` to handle cases where HOME is wrong
        let fypp_dir = std::process::Command::new("sh")
            .args([
                "-c",
                "dirname $(which fypp 2>/dev/null) 2>/dev/null || echo ''",
            ])
            .output()
            .ok()
            .and_then(|o| String::from_utf8(o.stdout).ok())
            .map(|s| s.trim().to_string())
            .filter(|s| !s.is_empty());
        let base_path = match env::var("PATH") {
            Ok(p) => p,
            Err(_) => "/usr/local/bin:/usr/bin:/bin".to_string(),
        };
        match fypp_dir {
            Some(dir) if dir != local_bin => format!("{local_bin}:{dir}:{base_path}"),
            _ => format!("{local_bin}:{base_path}"),
        }
    };

    // Skip cmake configure if already done (DBCSRConfig.cmake present)
    let dbcsr_config = dbcsr_install_dir
        .join("lib")
        .join("cmake")
        .join("dbcsr")
        .join("DBCSRConfig.cmake");
    if !dbcsr_config.exists() {
        println!("cargo:warning=Configuring DBCSR with CMake...");
        let num_jobs = env::var("NUM_JOBS").unwrap_or_else(|_| num_cpus::get().to_string());
        let configure = Command::new("cmake")
            .args([
                "-S",
                dbcsr_src_dir.to_str().unwrap(),
                "-B",
                dbcsr_build_dir.to_str().unwrap(),
                "-DCMAKE_BUILD_TYPE=RelWithDebInfo",
                "-DBUILD_SHARED_LIBS=OFF",
                "-DCMAKE_POSITION_INDEPENDENT_CODE=ON",
                "-DUSE_MPI=ON",
                "-DUSE_OPENMP=ON",
                "-DBUILD_TESTING=NO",
                "-DWITH_EXAMPLES=NO",
                "-DUSE_SMM=blas",
                "-DCMAKE_INSTALL_LIBDIR=lib",
                &format!("-DCMAKE_INSTALL_PREFIX={}", dbcsr_install_dir.display()),
                // Direct DBCSR to the statically-built OpenBLAS via explicit lib path
                &format!("-DBLAS_LIBRARIES={}", openblas_static_lib.display()),
                &format!("-DLAPACK_LIBRARIES={}", openblas_static_lib.display()),
                "-DCMAKE_PREFIX_PATH=/usr/lib64/openmpi;/usr/lib64",
                "-DCMAKE_Fortran_COMPILER=mpifort",
                "-DCMAKE_C_COMPILER=mpicc",
                "-DCMAKE_CXX_COMPILER=mpicxx",
                // Suppress -Werror=cast-function-type from OpenMPI C++ headers in DBCSR C++ examples
                &format!(
                    "-DCMAKE_CXX_FLAGS=-Wno-cast-function-type {}",
                    march_flag(&arch)
                ),
                &format!("-DCMAKE_C_FLAGS={}", march_flag(&arch)),
                &format!("-DCMAKE_Fortran_FLAGS={}", march_flag(&arch)),
                &format!("-DCMAKE_BUILD_PARALLEL_LEVEL={num_jobs}"),
            ])
            .env("PKG_CONFIG_PATH", &pkg_config_path)
            .env("PATH", &extended_path)
            .env("PYTHONPATH", &extended_pythonpath)
            .output()
            .expect("Failed to run cmake for DBCSR");
        if !configure.status.success() {
            eprintln!("DBCSR CMake configure failed!");
            eprintln!("stdout: {}", String::from_utf8_lossy(&configure.stdout));
            eprintln!("stderr: {}", String::from_utf8_lossy(&configure.stderr));
            panic!("DBCSR CMake configuration failed");
        }

        println!("cargo:warning=Building DBCSR...");
        let num_jobs = env::var("NUM_JOBS").unwrap_or_else(|_| num_cpus::get().to_string());
        let build = Command::new("cmake")
            .args([
                "--build",
                dbcsr_build_dir.to_str().unwrap(),
                "-j",
                &num_jobs,
            ])
            .env("PATH", &extended_path)
            .env("PYTHONPATH", &extended_pythonpath)
            .output()
            .expect("Failed to build DBCSR");
        if !build.status.success() {
            eprintln!("DBCSR build failed!");
            eprintln!("stdout: {}", String::from_utf8_lossy(&build.stdout));
            eprintln!("stderr: {}", String::from_utf8_lossy(&build.stderr));
            panic!("DBCSR build failed");
        }

        println!("cargo:warning=Installing DBCSR...");
        let install = Command::new("cmake")
            .args(["--install", dbcsr_build_dir.to_str().unwrap()])
            .env("PATH", &extended_path)
            .output()
            .expect("Failed to install DBCSR");
        if !install.status.success() {
            eprintln!("DBCSR install failed!");
            eprintln!("stdout: {}", String::from_utf8_lossy(&install.stdout));
            eprintln!("stderr: {}", String::from_utf8_lossy(&install.stderr));
            panic!("DBCSR install failed");
        }
        println!(
            "cargo:warning=DBCSR installed to {}",
            dbcsr_install_dir.display()
        );
    } else {
        println!(
            "cargo:warning=DBCSR already installed at {}",
            dbcsr_install_dir.display()
        );
    }

    dbcsr_install_dir
}

/// Download, build, and install FFTW3 (double precision, OpenMP variant) from the release tarball.
/// We use the release tarball (not the git repo) because the git repo does not include the
/// pre-generated codelet source files; those are generated by the OCaml `genfft` tool and
/// bundled only in official release archives.  The Autotools build correctly compiles and
/// links all solvtab/codelet objects that `libfftw3.a` requires internally.
/// Returns the install prefix directory.
fn build_fftw3(out_dir: &Path) -> PathBuf {
    let profile_dir = out_dir.ancestors().nth(3).unwrap_or(out_dir).to_path_buf();
    let arch = cp2k_target_arch();
    let stable_dir = profile_dir.join(format!("fftw3-stable{}", arch_dir_suffix(&arch)));
    std::fs::create_dir_all(&stable_dir).expect("Failed to create fftw3-stable dir");

    let fftw3_src_dir = stable_dir.join("src");
    let fftw3_build_dir = stable_dir.join("build");
    let fftw3_install_dir = stable_dir.join("install");

    // Skip entire build if already installed
    let fftw3_omp_lib = fftw3_install_dir.join("lib").join("libfftw3_omp.a");
    if fftw3_omp_lib.exists() {
        println!(
            "cargo:warning=FFTW3 already installed at {}",
            fftw3_install_dir.display()
        );
        return fftw3_install_dir;
    }

    // Download the FFTW3 release tarball if not present.  The release tarball includes the
    // pre-generated codelet files (dft/codelets/standard/, rdft/codelets/*, etc.) that are
    // required for a complete build with all solvtab symbols.
    if !fftw3_src_dir.exists() {
        println!("cargo:warning=Downloading FFTW3 3.3.10 release tarball...");
        let tarball = stable_dir.join("fftw-3.3.10.tar.gz");
        if !tarball.exists() {
            let dl = Command::new("curl")
                .args([
                    "-L",
                    "--retry",
                    "3",
                    "-o",
                    tarball.to_str().unwrap(),
                    "https://www.fftw.org/fftw-3.3.10.tar.gz",
                ])
                .output()
                .expect("Failed to download FFTW3 (is curl installed?)");
            if !dl.status.success() {
                panic!(
                    "Failed to download FFTW3 tarball: {}",
                    String::from_utf8_lossy(&dl.stderr)
                );
            }
        }
        std::fs::create_dir_all(&fftw3_src_dir).expect("Failed to create fftw3 src dir");
        let extract = Command::new("tar")
            .args([
                "-xzf",
                tarball.to_str().unwrap(),
                "--strip-components=1",
                "-C",
                fftw3_src_dir.to_str().unwrap(),
            ])
            .output()
            .expect("Failed to extract FFTW3 tarball");
        if !extract.status.success() {
            panic!(
                "Failed to extract FFTW3: {}",
                String::from_utf8_lossy(&extract.stderr)
            );
        }
        println!("cargo:warning=FFTW3 source extracted successfully");
    } else {
        println!("cargo:warning=FFTW3 source already exists, reusing");
    }

    std::fs::create_dir_all(&fftw3_build_dir).expect("Failed to create FFTW3 build dir");

    let num_jobs = env::var("NUM_JOBS").unwrap_or_else(|_| num_cpus::get().to_string());

    // Out-of-source autotools build.  --enable-openmp produces libfftw3_omp.a in addition
    // to the main libfftw3.a.  The Fortran interface (dfftw_* symbols) must NOT be disabled
    // because CP2K's fftw3_lib.F uses it.  These are just C wrappers, not Fortran code.
    println!("cargo:warning=Configuring FFTW3 with autotools...");
    let mut fftw3_args: Vec<String> = vec![
        format!("--prefix={}", fftw3_install_dir.display()),
        "--disable-shared".to_string(),
        "--enable-static".to_string(),
        "--with-pic".to_string(),
        "--enable-openmp".to_string(),
        "--enable-threads".to_string(),
    ];
    for flag in fftw3_simd_args(&arch) {
        fftw3_args.push(flag.to_string());
    }
    let configure = Command::new("sh")
        .arg(fftw3_src_dir.join("configure"))
        .args(&fftw3_args)
        .env("CFLAGS", march_flag(&arch))
        .current_dir(&fftw3_build_dir)
        .output()
        .expect("Failed to configure FFTW3");
    if !configure.status.success() {
        eprintln!("FFTW3 configure failed!");
        eprintln!("stdout: {}", String::from_utf8_lossy(&configure.stdout));
        eprintln!("stderr: {}", String::from_utf8_lossy(&configure.stderr));
        panic!("FFTW3 autotools configuration failed");
    }

    println!("cargo:warning=Building FFTW3...");
    let build = Command::new("make")
        .args(["-j", &num_jobs])
        .current_dir(&fftw3_build_dir)
        .output()
        .expect("Failed to build FFTW3");
    if !build.status.success() {
        eprintln!("FFTW3 build failed!");
        eprintln!("stdout: {}", String::from_utf8_lossy(&build.stdout));
        eprintln!("stderr: {}", String::from_utf8_lossy(&build.stderr));
        panic!("FFTW3 build failed");
    }

    println!("cargo:warning=Installing FFTW3...");
    let install = Command::new("make")
        .args(["install"])
        .current_dir(&fftw3_build_dir)
        .output()
        .expect("Failed to install FFTW3");
    if !install.status.success() {
        eprintln!("FFTW3 install failed!");
        eprintln!("stdout: {}", String::from_utf8_lossy(&install.stdout));
        eprintln!("stderr: {}", String::from_utf8_lossy(&install.stderr));
        panic!("FFTW3 install failed");
    }

    println!(
        "cargo:warning=FFTW3 installed to {}",
        fftw3_install_dir.display()
    );
    fftw3_install_dir
}

/// Clone, build, and install Reference-ScaLAPACK statically against our static OpenBLAS.
/// Returns the install prefix directory.
fn build_scalapack(out_dir: &Path, openblas_static_lib: &Path) -> PathBuf {
    let profile_dir = out_dir.ancestors().nth(3).unwrap_or(out_dir).to_path_buf();
    let arch = cp2k_target_arch();
    let stable_dir = profile_dir.join(format!("scalapack-stable{}", arch_dir_suffix(&arch)));
    std::fs::create_dir_all(&stable_dir).expect("Failed to create scalapack-stable dir");

    let scalapack_src_dir = stable_dir.join("src");
    let scalapack_build_dir = stable_dir.join("build");
    let scalapack_install_dir = stable_dir.join("install");

    // Skip if already installed
    let scalapack_lib = scalapack_install_dir.join("lib").join("libscalapack.a");
    if scalapack_lib.exists() {
        println!(
            "cargo:warning=ScaLAPACK already installed at {}",
            scalapack_install_dir.display()
        );
        return scalapack_install_dir;
    }

    if !scalapack_src_dir.exists() {
        println!("cargo:warning=Cloning Reference-ScaLAPACK repository (v2.2.0)...");
        let output = Command::new("git")
            .args([
                "clone",
                "--depth",
                "1",
                "--branch",
                "v2.2.0",
                "https://github.com/Reference-ScaLAPACK/scalapack.git",
                scalapack_src_dir.to_str().unwrap(),
            ])
            .output()
            .expect("Failed to clone ScaLAPACK repository. Make sure git is installed.");
        if !output.status.success() {
            panic!(
                "Failed to clone ScaLAPACK: {}",
                String::from_utf8_lossy(&output.stderr)
            );
        }
        println!("cargo:warning=ScaLAPACK repository cloned successfully");
    } else {
        println!("cargo:warning=ScaLAPACK repository already exists, reusing");
    }

    std::fs::create_dir_all(&scalapack_build_dir).expect("Failed to create ScaLAPACK build dir");

    let num_jobs = env::var("NUM_JOBS").unwrap_or_else(|_| num_cpus::get().to_string());

    // ScaLAPACK v2.2.0 has cmake_minimum_required(VERSION 2.8) in several files, which
    // CMake ≥ 3.27 rejects.  Patch every affected file in-place before configuring.
    // Note: -DCMAKE_POLICY_VERSION_MINIMUM only exists in CMake ≥ 3.27 and does not
    // help for the inner sub-cmake invocations anyway.
    for rel in &[
        "CMakeLists.txt",
        "BLACS/INSTALL/CMakeLists.txt",
        "scalapack_build.cmake",
    ] {
        let path = scalapack_src_dir.join(rel);
        if !path.exists() {
            continue;
        }
        let content =
            std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("Failed to read {rel}: {e}"));
        let patched = content
            .lines()
            .map(|line| {
                if line
                    .trim_start()
                    .to_ascii_lowercase()
                    .starts_with("cmake_minimum_required")
                {
                    "cmake_minimum_required(VERSION 3.5)".to_string()
                } else {
                    line.to_string()
                }
            })
            .collect::<Vec<_>>()
            .join("\n")
            + "\n";
        std::fs::write(&path, patched).unwrap_or_else(|e| panic!("Failed to patch {rel}: {e}"));
    }

    println!("cargo:warning=Configuring ScaLAPACK with CMake...");
    let configure = Command::new("cmake")
        .args([
            "-S",
            scalapack_src_dir.to_str().unwrap(),
            "-B",
            scalapack_build_dir.to_str().unwrap(),
            "-DCMAKE_BUILD_TYPE=Release",
            "-DBUILD_SHARED_LIBS=OFF",
            "-DCMAKE_POSITION_INDEPENDENT_CODE=ON",
            // Skip test executables – they drag in extra BLAS link deps at configure time
            "-DBUILD_TESTING=OFF",
            "-DCMAKE_INSTALL_LIBDIR=lib",
            &format!("-DCMAKE_INSTALL_PREFIX={}", scalapack_install_dir.display()),
            // Point directly at our static OpenBLAS for BLAS + LAPACK
            &format!("-DBLAS_LIBRARIES={}", openblas_static_lib.display()),
            &format!("-DLAPACK_LIBRARIES={}", openblas_static_lib.display()),
            "-DCMAKE_Fortran_COMPILER=mpifort",
            "-DCMAKE_C_COMPILER=mpicc",
            // pstrsen.f / pstrord.f in v2.2.0 pass scalars where arrays are expected;
            // gfortran >=10 treats this as an error by default.
            &format!(
                "-DCMAKE_Fortran_FLAGS=-fallow-argument-mismatch {}",
                march_flag(&arch)
            ),
            // igsum2d_.c et al. call BI_imvcopy/BI_svmcopy without a prior declaration;
            // GCC >=14 promotes -Wimplicit-function-declaration to an error by default.
            &format!(
                "-DCMAKE_C_FLAGS=-Wno-implicit-function-declaration {}",
                march_flag(&arch)
            ),
            &format!("-DCMAKE_BUILD_PARALLEL_LEVEL={num_jobs}"),
        ])
        .output()
        .expect("Failed to run cmake for ScaLAPACK");
    if !configure.status.success() {
        eprintln!("ScaLAPACK CMake configure failed!");
        eprintln!("stdout: {}", String::from_utf8_lossy(&configure.stdout));
        eprintln!("stderr: {}", String::from_utf8_lossy(&configure.stderr));
        panic!("ScaLAPACK CMake configuration failed");
    }

    println!("cargo:warning=Building ScaLAPACK...");
    let build = Command::new("cmake")
        .args([
            "--build",
            scalapack_build_dir.to_str().unwrap(),
            // Build only the library target.  Reference-ScaLAPACK v2.2.0 adds
            // test executables unconditionally (ignoring -DBUILD_TESTING=OFF),
            // so the default "all" target tries to link them against OpenBLAS
            // but without -lgomp, which fails when OpenBLAS is OpenMP-enabled.
            "--target",
            "scalapack",
            "-j",
            &num_jobs,
        ])
        .output()
        .expect("Failed to build ScaLAPACK");
    if !build.status.success() {
        eprintln!("ScaLAPACK build failed!");
        eprintln!("stdout: {}", String::from_utf8_lossy(&build.stdout));
        eprintln!("stderr: {}", String::from_utf8_lossy(&build.stderr));
        panic!("ScaLAPACK build failed");
    }

    println!("cargo:warning=Installing ScaLAPACK...");
    let install = Command::new("cmake")
        .args(["--install", scalapack_build_dir.to_str().unwrap()])
        .output()
        .expect("Failed to install ScaLAPACK");
    if !install.status.success() {
        eprintln!("ScaLAPACK install failed!");
        eprintln!("stdout: {}", String::from_utf8_lossy(&install.stdout));
        eprintln!("stderr: {}", String::from_utf8_lossy(&install.stderr));
        panic!("ScaLAPACK install failed");
    }

    println!(
        "cargo:warning=ScaLAPACK installed to {}",
        scalapack_install_dir.display()
    );
    scalapack_install_dir
}

fn configure_cp2k_cmake(
    cp2k_src_dir: &Path,
    cp2k_build_dir: &Path,
    dbcsr_install_dir: &Path,
    openblas_static_lib: &Path,
    openblas_include_dir: &Path,
    fftw3_install: &Path,
    scalapack_install: &Path,
) {
    println!("cargo:warning=Configuring CP2K with CMake...");

    let num_jobs = env::var("NUM_JOBS").unwrap_or_else(|_| num_cpus::get().to_string());

    // Build up PKG_CONFIG_PATH: prepend the OpenMPI pkgconfig dir so CMake's
    // FindSCALAPACK can locate scalapack.pc via pkg-config.
    let openmpi_pkgconfig = "/usr/lib64/openmpi/lib/pkgconfig";
    let pkg_config_path = match env::var("PKG_CONFIG_PATH") {
        Ok(existing) => format!("{openmpi_pkgconfig}:{existing}"),
        Err(_) => openmpi_pkgconfig.to_string(),
    };

    // CMAKE_PREFIX_PATH: FFTW3, DBCSR, and ScaLAPACK install dirs + system paths for MPI.
    // OpenBLAS is passed via explicit library/include variables.
    // ScaLAPACK v2.2.0 installs SCALAPACKConfig.cmake under its prefix, so adding it here
    // lets CP2K's find_package(SCALAPACK) discover the static library automatically.
    let cmake_prefix_path = format!(
        "{};{};{};/usr/lib64/openmpi;/usr/lib64",
        fftw3_install.display(),
        dbcsr_install_dir.display(),
        scalapack_install.display()
    );

    let output = Command::new("cmake")
        .args([
            "-S",
            cp2k_src_dir.to_str().unwrap(),
            "-B",
            cp2k_build_dir.to_str().unwrap(),
            "-DCMAKE_BUILD_TYPE=Release",
            "-DBUILD_SHARED_LIBS=OFF",
            "-DCP2K_USE_MPI=ON",
            "-DCP2K_USE_FFTW3=ON",
            "-DCP2K_USE_BLAS=ON",
            "-DCP2K_USE_LAPACK=ON",
            "-DCP2K_USE_SCALAPACK=ON",
            "-DCP2K_USE_LIBXC=OFF",
            "-DCP2K_USE_LIBINT2=OFF",
            "-DCMAKE_POSITION_INDEPENDENT_CODE=ON",
            // Use statically-built OpenBLAS for BLAS + LAPACK.
            // CUSTOM vendor lets us bypass auto-detection entirely.
            "-DCP2K_BLAS_VENDOR=CUSTOM",
            &format!(
                "-DCP2K_BLAS_LINK_LIBRARIES={}",
                openblas_static_lib.display()
            ),
            &format!(
                "-DCP2K_LAPACK_LINK_LIBRARIES={}",
                openblas_static_lib.display()
            ),
            &format!(
                "-DCP2K_BLAS_INCLUDE_DIRS={}",
                openblas_include_dir.display()
            ),
            // Use statically-built FFTW3 with OpenMP support
            &format!("-DCP2K_FFTW3_ROOT={}", fftw3_install.display()),
            "-DCP2K_ENABLE_FFTW3_OPENMP_SUPPORT=ON",
            &format!("-DCMAKE_PREFIX_PATH={cmake_prefix_path}"),
            "-DCMAKE_Fortran_COMPILER=mpifort",
            "-DCMAKE_C_COMPILER=mpicc",
            "-DCMAKE_CXX_COMPILER=mpicxx",
            &format!("-DCMAKE_C_FLAGS={}", march_flag(&cp2k_target_arch())),
            &format!("-DCMAKE_CXX_FLAGS={}", march_flag(&cp2k_target_arch())),
            &format!("-DCMAKE_Fortran_FLAGS={}", march_flag(&cp2k_target_arch())),
            &format!("-DCMAKE_BUILD_PARALLEL_LEVEL={num_jobs}"),
        ])
        .env("PKG_CONFIG_PATH", &pkg_config_path)
        .output()
        .expect("Failed to run cmake configure. Make sure cmake is installed.");

    if !output.status.success() {
        eprintln!("CMake configure failed!");
        eprintln!("stdout: {}", String::from_utf8_lossy(&output.stdout));
        eprintln!("stderr: {}", String::from_utf8_lossy(&output.stderr));
        panic!("CP2K CMake configuration failed");
    }

    println!("cargo:warning=CP2K CMake configuration completed");
}

fn build_cp2k_cmake(cp2k_build_dir: &Path) {
    println!("cargo:warning=Building CP2K with CMake - this may take 30+ minutes...");

    let num_jobs = env::var("NUM_JOBS").unwrap_or_else(|_| num_cpus::get().to_string());

    // Extend PATH and PYTHONPATH so fypp (used to preprocess Fortran files) is found.
    let home = env::var("HOME").unwrap_or_else(|_| "/root".to_string());
    let local_bin = format!("{home}/.local/bin");
    let extended_path = match env::var("PATH") {
        Ok(p) => format!("{local_bin}:{p}"),
        Err(_) => format!("{local_bin}:/usr/local/bin:/usr/bin:/bin"),
    };
    let extended_pythonpath = {
        let base = env::var("PYTHONPATH").unwrap_or_default();
        let candidates = [
            format!("{home}/.local/lib/python3.12/site-packages"),
            format!("{home}/.local/lib/python3.11/site-packages"),
            format!("{home}/.local/lib/python3.10/site-packages"),
            format!("{home}/.local/lib/python3.9/site-packages"),
        ];
        let extra = candidates.join(":");
        if base.is_empty() {
            extra
        } else {
            format!("{extra}:{base}")
        }
    };

    let output = Command::new("cmake")
        .args([
            "--build",
            cp2k_build_dir.to_str().unwrap(),
            "--target",
            "cp2k",
            "-j",
            &num_jobs,
        ])
        .env("PATH", &extended_path)
        .env("PYTHONPATH", &extended_pythonpath)
        .output()
        .expect("Failed to execute CP2K CMake build");

    if !output.status.success() {
        eprintln!("CP2K build failed!");
        eprintln!("stdout: {}", String::from_utf8_lossy(&output.stdout));
        eprintln!("stderr: {}", String::from_utf8_lossy(&output.stderr));
        panic!("CP2K compilation failed");
    }

    println!("cargo:warning=CP2K build completed successfully");
}

fn find_libcp2k(cp2k_build_dir: &Path) -> PathBuf {
    // CMake puts static libraries in {build_dir}/lib/
    let lib_path = cp2k_build_dir.join("lib").join("libcp2k.a");

    if lib_path.exists() {
        return lib_path;
    }

    panic!("Could not find libcp2k.a at {lib_path:?} after building");
}

fn configure_linking_without_cp2k(mpi_base_dir: &Path) {
    // Add MPI library directory
    if let Some(lib_dir) = find_mpi_lib_dir(mpi_base_dir) {
        println!("cargo:rustc-link-search=native={}", lib_dir.display());
    }

    // Link with required MPI libraries
    println!("cargo:rustc-link-lib=dylib=mpi");
    println!("cargo:rustc-link-lib=dylib=mpi_mpifh");
    println!("cargo:rustc-link-lib=dylib=gfortran");
    println!("cargo:rustc-link-lib=dylib=gomp");

    // Note: CP2K library linking will be handled separately when needed
    println!("cargo:warning=CP2K library linking skipped - header-only mode");
}

fn configure_linking(
    _lib_dir: &Path,
    mpi_base_dir: &Path,
    fftw3_install: &Path,
    openblas_static_lib: &Path,
    scalapack_install: &Path,
) {
    // Add MPI library directory
    if let Some(mpi_lib_dir) = find_mpi_lib_dir(mpi_base_dir) {
        println!("cargo:rustc-link-search=native={}", mpi_lib_dir.display());
    }

    // Add standard system library directories
    println!("cargo:rustc-link-search=native=/lib64");
    println!("cargo:rustc-link-search=native=/usr/lib64");
    println!("cargo:rustc-link-search=native=/usr/lib64/openmpi/lib");

    // Static FFTW3 (built from source by build_fftw3()).
    // Use --whole-archive so all codelet solvtab symbols are included;
    // FFTW3's conf.c references hundreds of solvtab symbols spread across
    // many object files and a single-pass linker would GC them otherwise.
    let fftw3_lib_dir = fftw3_install.join("lib");
    println!("cargo:rustc-link-search=native={}", fftw3_lib_dir.display());
    println!("cargo:rustc-link-arg=-Wl,--whole-archive");
    println!(
        "cargo:rustc-link-arg={}",
        fftw3_lib_dir.join("libfftw3_omp.a").display()
    );
    println!(
        "cargo:rustc-link-arg={}",
        fftw3_lib_dir.join("libfftw3.a").display()
    );
    println!("cargo:rustc-link-arg=-Wl,--no-whole-archive");

    // Static ScaLAPACK (built from source by build_scalapack()).
    // BLACS is bundled inside Reference-ScaLAPACK v2.2.0 — no separate library needed.
    let scalapack_lib_dir = scalapack_install.join("lib");
    println!(
        "cargo:rustc-link-search=native={}",
        scalapack_lib_dir.display()
    );
    println!("cargo:rustc-link-lib=static=scalapack");

    // Link with required libraries
    // Note: BLAS + LAPACK are provided as static=openblas by the openblas-src build dependency
    println!("cargo:rustc-link-lib=dylib=mpi");
    println!("cargo:rustc-link-lib=dylib=mpi_mpifh");
    println!("cargo:rustc-link-lib=dylib=gfortran");
    println!("cargo:rustc-link-lib=dylib=m");
    println!("cargo:rustc-link-lib=dylib=z");
    println!("cargo:rustc-link-lib=dylib=dl");
    println!("cargo:rustc-link-lib=dylib=pthread");
    println!("cargo:rustc-link-lib=dylib=gomp"); // OpenMP

    // Force OpenMP linking for Python extensions by adding link args
    println!("cargo:rustc-link-arg=-Wl,-rpath,/lib64");
    println!("cargo:rustc-link-arg=-lgomp");
    println!("cargo:rustc-link-arg=-lgfortran");

    // openblas-src emits cargo:rustc-link-lib=static=openblas early in the link
    // order (before libcp2k.a / libdbcsr.a).  With single-pass linkers, BLAS
    // symbols referenced by CP2K/DBCSR are not yet live when libopenblas.a is
    // first processed, so nothing gets pulled from the archive.  Adding the
    // archive again by its full path at the end of the link command guarantees
    // it is scanned after all dependent archives.
    println!("cargo:rustc-link-arg={}", openblas_static_lib.display());
}

#[cfg(feature = "bindgen")]
fn generate_bindings(manifest_dir: &str, out_dir: &Path) {
    let bindings = bindgen::Builder::default()
        .header(format!("{manifest_dir}/libcp2k.h"))
        .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
        .generate()
        .expect("Unable to generate bindings");

    bindings
        .write_to_file(out_dir.join("bindings.rs"))
        .expect("Couldn't write bindings!");
}

fn clone_cp2k_repository(out_dir: &Path) -> PathBuf {
    let cp2k_src_dir = out_dir.join("cp2k");

    if !cp2k_src_dir.exists() {
        println!("cargo:warning=Cloning CP2K repository...");

        let output = Command::new("git")
            .args([
                "clone",
                "--depth",
                "1",
                "--branch",
                "support/v2026.1", // Use a stable branch
                "--recursive",     // Clone submodules too
                "https://github.com/cp2k/cp2k.git",
                cp2k_src_dir.to_str().unwrap(),
            ])
            .output()
            .expect("Failed to clone CP2K repository. Make sure git is installed.");

        if !output.status.success() {
            panic!(
                "Failed to clone CP2K repository: {}",
                String::from_utf8_lossy(&output.stderr)
            );
        }

        println!("cargo:warning=CP2K repository cloned successfully");
    } else {
        println!("cargo:warning=CP2K repository already exists, updating submodules...");

        // Update submodules if the repository already exists
        let output = Command::new("git")
            .args(["submodule", "update", "--init", "--recursive"])
            .current_dir(&cp2k_src_dir)
            .output()
            .expect("Failed to update submodules");

        if !output.status.success() {
            panic!(
                "Failed to update submodules: {}",
                String::from_utf8_lossy(&output.stderr)
            );
        }
    }

    cp2k_src_dir
}

/// Build the Fortran extension against pre-existing CP2K mod files.
/// Used when 'extended' feature is enabled but 'build-cp2k' is not:
/// CP2K was built previously and mod files already exist.
fn build_fortran_extensions_from_mods(out_dir: &Path, _profile: &str) {
    let extensions_dir =
        PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()).join("extensions/fortran");
    let obj_dir = out_dir.join("extensions_obj");
    let lib_dir = out_dir.join("extensions_lib");

    std::fs::create_dir_all(&obj_dir).expect("Failed to create obj dir");
    std::fs::create_dir_all(&lib_dir).expect("Failed to create lib dir");

    // Locate pre-built CP2K mod files under target/<profile>/cp2k-stable.../build/src/mod_files
    let target_dir = out_dir.ancestors().nth(3).unwrap_or(out_dir);
    let arch = cp2k_target_arch();
    let mod_files_dir = target_dir
        .join(format!("cp2k-stable{}", arch_dir_suffix(&arch)))
        .join("build")
        .join("src")
        .join("mod_files");

    if !mod_files_dir.exists() {
        println!(
            "cargo:warning=CP2K mod files not found at {} — skipping extension build",
            mod_files_dir.display()
        );
        return;
    }

    // CP2K source directory for base includes
    let cp2k_src_dir = target_dir
        .join(format!("cp2k-stable{}", arch_dir_suffix(&arch)))
        .join("cp2k")
        .join("src");

    println!("cargo:warning=Compiling Fortran extensions against pre-built CP2K...");

    let inc_cp2k = format!("-I{}", cp2k_src_dir.display());
    let inc_base = format!("-I{}/base", cp2k_src_dir.display());
    let inc_mods = format!("-I{}", mod_files_dir.display());

    let fortran_files = vec!["libcp2k_extended.F90", "libcp2k_mpi.F90"];
    let mut object_files = Vec::new();

    for fortran_file in &fortran_files {
        let src_path = extensions_dir.join(fortran_file);
        if !src_path.exists() {
            println!("cargo:warning=Skipping {} — file not found", fortran_file);
            continue;
        }

        let obj_name = fortran_file.replace(".F90", ".o").replace(".f90", ".o");
        let obj_path = obj_dir.join(&obj_name);

        println!("cargo:warning=Compiling {}...", fortran_file);

        let args = vec![
            "-c",
            "-fPIC",
            "-O2",
            "-g",
            "-ffree-form",
            "-ffree-line-length-none",
            "-fallow-argument-mismatch",
            "-D__parallel",
            inc_cp2k.as_str(),
            inc_base.as_str(),
            inc_mods.as_str(),
            "-o",
            obj_path.to_str().unwrap(),
            src_path.to_str().unwrap(),
        ];

        let output = Command::new("mpifort")
            .args(&args)
            .output()
            .unwrap_or_else(|_| panic!("Failed to invoke mpifort for {}", fortran_file));

        if !output.status.success() {
            eprintln!("\n=== Fortran compilation failed for {} ===", fortran_file);
            eprintln!("Command: mpifort {}", args.join(" "));
            eprintln!("stdout: {}", String::from_utf8_lossy(&output.stdout));
            eprintln!("stderr: {}", String::from_utf8_lossy(&output.stderr));
            panic!("Fortran compilation failed for {}", fortran_file);
        }

        object_files.push(obj_path);
    }

    if object_files.is_empty() {
        println!("cargo:warning=No Fortran extension files compiled");
        return;
    }

    // Create static library
    let lib_path = lib_dir.join("libcp2k_extended.a");
    println!("cargo:warning=Creating libcp2k_extended.a...");

    let output = Command::new("ar")
        .arg("rcs")
        .arg(&lib_path)
        .args(&object_files)
        .output()
        .expect("Failed to create archive");

    if !output.status.success() {
        eprintln!("stdout: {}", String::from_utf8_lossy(&output.stdout));
        eprintln!("stderr: {}", String::from_utf8_lossy(&output.stderr));
        panic!("Failed to create libcp2k_extended.a");
    }

    // Link the extended library
    println!("cargo:rustc-link-search=native={}", lib_dir.display());
    println!("cargo:rustc-link-lib=static=cp2k_extended");
    println!(
        "cargo:warning=Fortran extensions built successfully at {}",
        lib_path.display()
    );
}

fn build_fortran_extensions(cp2k_src_dir: &Path, cp2k_build_dir: &Path, out_dir: &Path) {
    let extensions_dir =
        PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()).join("extensions/fortran");
    let obj_dir = out_dir.join("extensions_obj");
    let lib_dir = out_dir.join("extensions_lib");

    std::fs::create_dir_all(&obj_dir).expect("Failed to create obj dir");
    std::fs::create_dir_all(&lib_dir).expect("Failed to create lib dir");

    println!("cargo:warning=Compiling Fortran extensions...");

    // Get the CP2K source directory for includes
    let cp2k_src = cp2k_src_dir.join("src");

    // In CMake builds, .mod files are placed in {build_dir}/src/mod_files/
    let mod_files_dir = cp2k_build_dir.join("src").join("mod_files");

    if !mod_files_dir.exists() {
        panic!(
            "CP2K mod files directory not found: {}. Make sure CP2K is built before compiling extensions.",
            mod_files_dir.display()
        );
    }

    // Compile each Fortran source file
    let fortran_files = vec!["libcp2k_extended.F90", "libcp2k_mpi.F90"];

    let mut object_files = Vec::new();

    for fortran_file in &fortran_files {
        let src_path = extensions_dir.join(fortran_file);
        if !src_path.exists() {
            println!("cargo:warning=Skipping {} - file not found", fortran_file);
            continue;
        }

        let obj_name = fortran_file.replace(".F90", ".o").replace(".f90", ".o");
        let obj_path = obj_dir.join(&obj_name);

        println!("cargo:warning=Compiling {}...", fortran_file);

        let output = Command::new("mpifort")
            .args([
                "-c",
                "-fPIC",
                "-O2",
                "-g",
                "-ffree-form",
                "-ffree-line-length-none",
                "-fallow-argument-mismatch",
                "-D__parallel", // Enable MPI code paths in CP2K Fortran extensions
                &format!("-I{}", cp2k_src.display()),
                &format!("-I{}/base", cp2k_src.display()),
                &format!("-I{}", mod_files_dir.display()), // Add path to .mod files
                "-o",
                obj_path.to_str().unwrap(),
                src_path.to_str().unwrap(),
            ])
            .output()
            .unwrap_or_else(|_| panic!("Failed to compile {}", fortran_file));

        if !output.status.success() {
            eprintln!("Failed to compile {}", fortran_file);
            eprintln!("stdout: {}", String::from_utf8_lossy(&output.stdout));
            eprintln!("stderr: {}", String::from_utf8_lossy(&output.stderr));
            panic!("Fortran compilation failed for {}", fortran_file);
        }

        object_files.push(obj_path);
    }

    if object_files.is_empty() {
        println!("cargo:warning=No Fortran extension files compiled");
        return;
    }

    // Create a static library from the object files
    let lib_path = lib_dir.join("libcp2k_extended.a");

    println!("cargo:warning=Creating libcp2k_extended.a...");

    let output = Command::new("ar")
        .arg("rcs")
        .arg(&lib_path)
        .args(&object_files)
        .output()
        .expect("Failed to create archive");

    if !output.status.success() {
        eprintln!("stdout: {}", String::from_utf8_lossy(&output.stdout));
        eprintln!("stderr: {}", String::from_utf8_lossy(&output.stderr));
        panic!("Failed to create libcp2k_extended.a");
    }

    println!(
        "cargo:warning=Fortran extensions built successfully at {}",
        lib_path.display()
    );
}

/// Libraries to bundle next to the worker binary (version-sensitive, absent from
/// the manylinux allowlist, or required Fortran runtime).
/// Note: BLAS/LAPACK (libopenblas), FFTW3, and ScaLAPACK are now statically linked and
/// no longer appear as dynamic dependencies of the worker binary.
const BUNDLE_PATTERNS: &[&str] = &["libgfortran.so", "libquadmath.so", "libgcc_s.so"];

/// Libraries that must remain dynamic – MPI ABI is interconnect-specific; glibc is always present.
/// Note: ScaLAPACK is now statically linked and no longer needs to be kept dynamic.
const KEEP_DYNAMIC_PATTERNS: &[&str] = &[
    "libmpi",
    "libopen-",
    "libhwloc",
    "libevent",
    "libc.so",
    "libm.so",
    "libmvec",
    "libz.so",
    "libdl.so",
    "libpthread",
    "libgomp",
    "linux-vdso",
    "ld-linux",
];

fn should_bundle_lib(soname: &str) -> bool {
    for pat in KEEP_DYNAMIC_PATTERNS {
        if soname.contains(pat) {
            return false;
        }
    }
    for pat in BUNDLE_PATTERNS {
        if soname.contains(pat) {
            return true;
        }
    }
    false
}

/// Run `ldd` on `binary` and return (soname, resolved_path) pairs.
fn ldd_deps(binary: &Path) -> Vec<(String, PathBuf)> {
    let output = match Command::new("ldd").arg(binary).output() {
        Ok(o) if o.status.success() => o,
        _ => return vec![],
    };
    let stdout = String::from_utf8_lossy(&output.stdout);
    let mut result = Vec::new();
    for line in stdout.lines() {
        // Format:  libfoo.so.1 => /path/to/libfoo.so.1 (0x...)
        let parts: Vec<&str> = line.trim().splitn(4, ' ').collect();
        if parts.len() >= 3 && parts[1] == "=>" && !parts[2].is_empty() && parts[2] != "not" {
            let soname = parts[0].to_string();
            let resolved = PathBuf::from(parts[2]);
            if resolved.exists() {
                result.push((soname, resolved));
            }
        }
    }
    result
}

/// Copy the worker binary and bundleable shared libraries into `cp2k_rs/` so
/// maturin includes them in the wheel. Patches the worker RPATH with `$ORIGIN`
/// so bundled libs are found at runtime.
///
/// The worker binary must already exist at `target/<profile>/cp2k_rs_worker`.
/// Use `bundle-wheel.sh` (or build the binary manually first) before running
/// `maturin build`.
fn bundle_worker_libs(manifest_dir: &str, profile: &str) {
    let worker = PathBuf::from(manifest_dir)
        .join("target")
        .join(profile)
        .join("cp2k_rs_worker");

    if !worker.exists() {
        // Worker not built yet – skip silently. The bundle-wheel.sh script
        // ensures the binary is built before maturin build runs.
        println!(
            "cargo:warning=cp2k_rs_worker not found at {} – skipping lib bundling (run bundle-wheel.sh to build the wheel)",
            worker.display()
        );
        return;
    }

    let pkg_dir = PathBuf::from(manifest_dir).join("cp2k_rs");
    std::fs::create_dir_all(&pkg_dir).expect("Failed to create cp2k_rs dir");

    // Copy the worker binary itself (maturin packages it from cp2k_rs/)
    let dst_worker = pkg_dir.join("cp2k_rs_worker");
    let src_mtime = std::fs::metadata(&worker).and_then(|m| m.modified()).ok();
    let dst_mtime = std::fs::metadata(&dst_worker)
        .and_then(|m| m.modified())
        .ok();
    if !dst_worker.exists() || src_mtime != dst_mtime {
        std::fs::copy(&worker, &dst_worker).expect("Failed to copy cp2k_rs_worker");
        // Ensure executable bit
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let mut perms = std::fs::metadata(&dst_worker).unwrap().permissions();
            perms.set_mode(perms.mode() | 0o111);
            std::fs::set_permissions(&dst_worker, perms).unwrap();
        }
        println!("cargo:warning=Copied cp2k_rs_worker to cp2k_rs/");
    }

    // Find and copy bundleable libs
    let deps = ldd_deps(&dst_worker);
    let mut bundled: Vec<String> = Vec::new();

    for (soname, resolved) in &deps {
        if !should_bundle_lib(soname) {
            continue;
        }
        let dest = pkg_dir.join(resolved.file_name().unwrap());
        if !dest.exists() {
            std::fs::copy(resolved, &dest)
                .unwrap_or_else(|e| panic!("Failed to copy {}: {}", resolved.display(), e));
            // Make writable so patchelf can modify it if needed
            #[cfg(unix)]
            {
                use std::os::unix::fs::PermissionsExt;
                let mut perms = std::fs::metadata(&dest).unwrap().permissions();
                perms.set_mode(perms.mode() | 0o200);
                std::fs::set_permissions(&dest, perms).unwrap();
            }
            println!(
                "cargo:warning=Bundled {}",
                resolved.file_name().unwrap().to_string_lossy()
            );
        }
        bundled.push(resolved.file_name().unwrap().to_string_lossy().into_owned());
    }

    if bundled.is_empty() {
        return;
    }

    // Patch RPATH on the worker and all bundled libs to include $ORIGIN so
    // they find each other at runtime (e.g. libblas.so.3 finding libflexiblas.so.3).
    let patchelf = match which_patchelf() {
        Some(p) => p,
        None => {
            println!("cargo:warning=patchelf not found – bundled libs may not be found at runtime");
            return;
        }
    };

    let mut to_patch = vec![dst_worker.clone()];
    for name in &bundled {
        to_patch.push(pkg_dir.join(name));
    }

    for path in &to_patch {
        patch_rpath_origin(&patchelf, path);
    }
}

fn bundle_cp2k_data(manifest_dir: &str, profile: &str) {
    let profile_dir = PathBuf::from(manifest_dir).join("target").join(profile);

    let candidates = [
        profile_dir.join("cp2k-stable").join("cp2k").join("data"),
        profile_dir.join("cp2k-stable-v4").join("cp2k").join("data"),
    ];

    let src_dir = candidates.iter().find(|p| p.exists());
    let Some(src_dir) = src_dir else {
        println!(
            "cargo:warning=CP2K data directory not found under target/{profile}/cp2k-stable*/cp2k/data – skipping data bundling"
        );
        return;
    };

    let data_dir = PathBuf::from(manifest_dir).join("cp2k_rs").join("data");
    std::fs::create_dir_all(&data_dir).expect("Failed to create cp2k_rs/data");

    let files = ["BASIS_MOLOPT", "GTH_POTENTIALS"];
    for name in &files {
        let src_file = src_dir.join(name);
        let dst_file = data_dir.join(name);
        if src_file.exists() {
            std::fs::copy(&src_file, &dst_file).unwrap_or_else(|e| {
                panic!("Failed to copy CP2K data file {}: {e}", src_file.display())
            });
            println!("cargo:warning=Bundled CP2K data file {}", name);
        } else {
            println!(
                "cargo:warning=CP2K data file '{}' not found in {}",
                name,
                src_dir.display()
            );
        }
    }
}

fn patch_rpath_origin(patchelf: &Path, path: &Path) {
    let current_rpath = Command::new(patchelf)
        .args(["--print-rpath", path.to_str().unwrap()])
        .output()
        .ok()
        .filter(|o| o.status.success())
        .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string())
        .unwrap_or_default();

    if !current_rpath.contains("$ORIGIN") {
        let new_rpath = if current_rpath.is_empty() {
            "$ORIGIN".to_string()
        } else {
            format!("$ORIGIN:{current_rpath}")
        };
        let status = Command::new(patchelf)
            .args(["--set-rpath", &new_rpath, path.to_str().unwrap()])
            .status()
            .expect("Failed to run patchelf");
        if status.success() {
            println!(
                "cargo:warning=Set RPATH on {}: {new_rpath}",
                path.file_name().unwrap().to_string_lossy()
            );
        } else {
            println!(
                "cargo:warning=patchelf failed on {}",
                path.file_name().unwrap().to_string_lossy()
            );
        }
    }
}

fn which_patchelf() -> Option<PathBuf> {
    // Check PATH entries
    if let Ok(path_var) = env::var("PATH") {
        for dir in path_var.split(':') {
            let candidate = PathBuf::from(dir).join("patchelf");
            if candidate.exists() {
                return Some(candidate);
            }
        }
    }
    // Common extra locations
    for p in &["/usr/bin/patchelf", "/usr/local/bin/patchelf"] {
        let p = PathBuf::from(p);
        if p.exists() {
            return Some(p);
        }
    }
    None
}
