use std::env;
use std::path::PathBuf;

use bindgen::callbacks::ParseCallbacks;

struct Dirs {
    _manifest_dir: PathBuf,
    mldsa_src_dir: PathBuf,
    build_harness_dir: PathBuf,
    build_harness_extra_dir: PathBuf,
}

impl Dirs {
    fn new() -> Self {
        let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());

        let mldsa_src_dir = manifest_dir.join("mldsa-native").join("mldsa");
        let build_harness_dir = manifest_dir.join("mldsa-build-harness");
        let build_harness_extra_dir = manifest_dir.join("mldsa-build-harness-extra");

        Self {
            _manifest_dir: manifest_dir,
            mldsa_src_dir,
            build_harness_dir,
            build_harness_extra_dir,
        }
    }
}

#[cfg(feature = "native")]
mod build_native_support {
    use super::Dirs;
    use super::PathBuf;

    mod native_detect;

    pub fn configure_for_native(builder: &mut cc::Build, dirs: &Dirs) {
        println!("cargo:warning=(INFO) \"native\" feature flag enabled.");
        let detected = native_detect::detect();
        println!("cargo:warning=(INFO) {detected:?}");
        detected.apply(builder);

        builder
            // We provide our own `mld_sys_check_capability()` to dispatch
            // between native and portable implementations at runtime.
            .define("MLD_CONFIG_CUSTOM_CAPABILITY_FUNC", "")
            // Enables native arithmetic backend
            .define("MLD_CONFIG_USE_NATIVE_BACKEND_ARITH", "")
            // Enables native FIPS-202 backend
            .define("MLD_CONFIG_USE_NATIVE_BACKEND_FIPS202", "")
            // Adds the assembly sources as a separate compilation unit
            .file(dirs.build_harness_dir.join("mldsa_native_asm_all.S"));
    }

    #[derive(Debug)]
    struct StripEnumPrefix;

    impl bindgen::callbacks::ParseCallbacks for StripEnumPrefix {
        fn enum_variant_name(
            &self,
            enum_name: Option<&str>,
            original_variant_name: &str,
            _variant_value: bindgen::callbacks::EnumVariantValue,
        ) -> Option<String> {
            if Some("mld_sys_cap") == enum_name && original_variant_name.starts_with("MLD_SYS_CAP_")
            {
                // original_variant_name is the raw C name, e.g. "MLD_SYS_CAP_X86_64_AVX2"
                // We strip the prefix, because bindgen will prepend the `enum_name`.
                Some(
                    original_variant_name
                        .strip_prefix("MLD_SYS_CAP_")
                        .map(|s| s.to_string())?,
                )
            } else {
                None
            }
        }
    }

    pub fn generate_support_bindings(dirs: &Dirs) {
        let header_file = dirs.build_harness_dir.join("detect_capabilities.h");

        let bindings = bindgen::Builder::default()
            .clang_args([
                format!("-I{}", dirs.mldsa_src_dir.to_string_lossy()),
                format!("-I{}", dirs.build_harness_dir.to_string_lossy()),
                format!("-D{}", "MLD_CONFIG_CUSTOM_CAPABILITY_FUNC"),
            ])
            .header(header_file.to_string_lossy())
            // Use ctypes from ::core
            .use_core()
            // Use our custom parsing rules for renaming enums
            .parse_callbacks(Box::new(StripEnumPrefix))
            .constified_enum_module("mld_sys_cap")
            // Tell cargo to invalidate the built crate whenever any of the
            // included header files changed.
            .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
            // Finish the builder and generate the bindings.
            .generate()
            // Unwrap the Result and panic on failure.
            .expect("Unable to generate bindings");

        // Write the bindings to the $OUT_DIR/bindings.rs file.
        let out_path = PathBuf::from(std::env::var("OUT_DIR").unwrap());
        bindings
            .write_to_file(out_path.join("detect_capabilities_bindings.rs"))
            .expect("Couldn't write bindings!");
    }
}

fn compile_c_sources(dirs: &Dirs) {
    let cc_flags = [
        "-Wall",
        "-Wextra",
        "-Werror=unused-result",
        "-Wpedantic",
        "-Werror",
        "-Wmissing-prototypes",
        "-Wshadow",
        "-Wpointer-arith",
        "-Wredundant-decls",
        "-Wconversion",
        "-Wsign-conversion",
        "-Wno-long-long",
        "-Wno-unknown-pragmas",
        "-Wno-unused-command-line-argument",
        "-O3",
        "-fomit-frame-pointer",
        "-std=c99",
        "-pedantic",
        "-MMD",
    ];

    const RANDOMBYTES_INTERNAL_NAME: &str = "_mldsa_native_rs_internal_randombytes";
    println!("cargo:rustc-env=RANDOMBYTES_INTERNAL_NAME={RANDOMBYTES_INTERNAL_NAME}");

    let mut builder = cc::Build::new();

    builder
        .flags(&cc_flags)
        .define("randombytes", RANDOMBYTES_INTERNAL_NAME)
        .define("MLD_CONFIG_NAMESPACE_PREFIX", "mldsa")
        .includes([&dirs.mldsa_src_dir, &dirs.build_harness_dir])
        .file(dirs.build_harness_dir.join("mldsa_native_all.c"));

    #[cfg(feature = "native")]
    build_native_support::configure_for_native(&mut builder, dirs);
    #[cfg(feature = "native")]
    build_native_support::generate_support_bindings(dirs);

    builder.compile("mldsa_native");

    println!(
        "cargo::rerun-if-changed={}",
        dirs.mldsa_src_dir.to_string_lossy()
    );
    println!(
        "cargo::rerun-if-changed={}",
        dirs.build_harness_dir.to_string_lossy()
    );
}

/// Wrap the entire doc comment in a "```text" block.
///
/// Any existing "```" delimiters are removed; the content of those blocks
/// remains and its formatting is otherwise preserved.
#[derive(Debug)]
struct CommentFormattingEscaper;

impl ParseCallbacks for CommentFormattingEscaper {
    fn process_comment(&self, comment: &str) -> Option<String> {
        // get rid of lines containing triple-backtick delimiters, since we're about to add our own
        let comment = comment
            .lines()
            .filter(|line| !line.contains("```"))
            .collect::<Vec<_>>()
            .join("\n");

        Some(format!("```text\n{comment}\n```"))
    }
}

fn generate_bindings(dirs: &Dirs) {
    let wrapper_h_file = dirs.build_harness_extra_dir.join("wrapper.h");

    let bindings = bindgen::Builder::default()
        .clang_args([
            format!("-I{}", dirs.mldsa_src_dir.to_string_lossy()),
            format!("-I{}", dirs.build_harness_dir.to_string_lossy()),
        ])
        .use_core()
        .header(wrapper_h_file.to_string_lossy())
        // Tell cargo to invalidate the built crate whenever any of the
        // included header files changed.
        .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
        // Render all doc comments as preformatted plain text.
        .parse_callbacks(Box::new(CommentFormattingEscaper))
        // Finish the builder and generate the bindings.
        .generate()
        // Unwrap the Result and panic on failure.
        .expect("Unable to generate bindings");

    // Write the bindings to the $OUT_DIR/bindings.rs file.
    let out_path = PathBuf::from(std::env::var("OUT_DIR").unwrap());
    bindings
        .write_to_file(out_path.join("bindings.rs"))
        .expect("Couldn't write bindings!");
}

fn main() {
    let dirs = Dirs::new();

    compile_c_sources(&dirs);
    generate_bindings(&dirs);
}
