//
// hardened-malloc-sys: Rust bindings for GrapheneOS allocator
// build.rs: Helper file for build-time information
//
// Copyright (c) 2025 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: MIT

use std::{
    env, fs,
    io::Write,
    path::{Path, PathBuf},
};

fn read_config_file(path: &PathBuf) -> Vec<(String, String)> {
    let content =
        fs::read_to_string(path).expect(&format!("Failed to read config file {:?}", path));

    content
        .lines()
        .filter(|line| {
            let trimmed = line.trim();
            !trimmed.starts_with("#") && !trimmed.is_empty()
        })
        .filter_map(|line| {
            let parts: Vec<&str> = line.splitn(2, '=').collect();
            if parts.len() == 2 {
                Some((parts[0].trim().to_string(), parts[1].trim().to_string()))
            } else {
                panic!("Invalid config line: {line}!");
            }
        })
        .collect()
}

fn main() {
    // Path to the config directory (relative to the project root).
    let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));

    // Determine the config file based on the "light" feature.
    let config_file = if env::var_os("CARGO_FEATURE_LIGHT").is_some() {
        root.join("config-light.conf")
    } else {
        root.join("config-default.conf")
    };

    // Determine CONFIG_SEAL_METADATA based on "seal" feature and target_env=gnu.
    let feature_seal = cfg!(target_env = "gnu") && env::var_os("CARGO_FEATURE_SEAL").is_some();
    let seal_metadata = if feature_seal { Some("true") } else { None };

    // Determine CONFIG_CLASS_REGION_SIZE based on "small" and "tiny" features.
    let feature_small = env::var_os("CARGO_FEATURE_SMALL").is_some();
    let feature_tiny = env::var_os("CARGO_FEATURE_TINY").is_some();
    if feature_small && feature_tiny {
        // Prevent nonsensical use.
        panic!("At most one of small and tiny features must be specified!");
    }
    let class_region_size = if feature_small {
        Some(4294967296usize) // 4GiB
    } else if feature_tiny {
        Some(4194304usize) // 4MiB
    } else {
        None
    };

    // Check for C17 compliant compiler, panic on error.
    let out_dir = PathBuf::from(env::var_os("OUT_DIR").unwrap());
    check_compiler_c17(&out_dir);

    // Initialize the cc build system.
    let mut build = cc::Build::new();
    build.flag("-std=c17");
    build.define("_GNU_SOURCE", Some("1"));

    // Read configuration key-value pairs from the selected config file.
    // Add each config as a preprocessor define (-Dkey=value).
    let config = read_config_file(&config_file);
    for (key, value) in config {
        if key == "CONFIG_SEAL_METADATA" {
            // Override CONFIG_SEAL_METADATA based on "seal" feature and target_env=gnu.
            if let Some(val) = seal_metadata {
                build.define(&key, Some(val.to_string().as_str()));
                continue;
            }
        } else if key == "CONFIG_CLASS_REGION_SIZE" {
            // Override CONFIG_CLASS_REGION_SIZE based on "small" feature.
            if let Some(val) = class_region_size {
                build.define(&key, Some(val.to_string().as_str()));
                continue;
            }
        }
        build.define(&key, Some(value.as_str()));
    }

    // Set source files directory (vendor/hardened-malloc).
    let vendor_dir = root.join("vendor").join("hardened-malloc");

    // Set source files.
    build.file(vendor_dir.join("chacha.c"));
    build.file(vendor_dir.join("h_malloc.c"));
    build.file(vendor_dir.join("memory.c"));
    build.file(vendor_dir.join("pages.c"));
    build.file(vendor_dir.join("random.c"));
    build.file(vendor_dir.join("util.c"));

    // Set include files.
    build.include(&vendor_dir);
    build.include(vendor_dir.join("include"));
    build.include(vendor_dir.join("third_party"));

    // Add LDFLAGS equivalent for linker flags.
    println!("cargo:rustc-link-arg=-Wl,-O1");
    println!("cargo:rustc-link-arg=-Wl,--as-needed");
    println!("cargo:rustc-link-arg=-Wl,-z,defs");
    println!("cargo:rustc-link-arg=-Wl,-z,relro");
    println!("cargo:rustc-link-arg=-Wl,-z,now");
    println!("cargo:rustc-link-arg=-Wl,-z,nodlopen");
    println!("cargo:rustc-link-arg=-Wl,-z,text");

    // Compile.
    build.compile("hardened_malloc");

    // Link statically.
    println!("cargo:rustc-link-lib=static=hardened_malloc");
    println!("cargo:rustc-link-search={}", out_dir.display());

    // Rerun the build script if config files change.
    println!("cargo:rerun-if-changed={}", config_file.display());
    println!("cargo:rerun-if-changed={}", vendor_dir.display());
}

// Check for C17 compliant compiler, panic on error.
fn check_compiler_c17<P: AsRef<Path>>(out_dir: P) {
    // Write test file.
    let test = PathBuf::from(out_dir.as_ref()).join("test_c17.c");
    let mut file = fs::File::create(&test).unwrap();
    writeln!(file, "int main() {{ return 0; }}").unwrap();

    // Initialize the cc build system.
    let mut build = cc::Build::new();
    build.file(test);
    build.flag("-std=c17");

    // Attempt to compile the test file.
    if let Err(error) = build.try_compile("test_c17") {
        panic!("hardened-malloc-sys requires a C17 supporting compiler: {error}!");
    }
}
