// Copyright (C) Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: Apache-2.0

// See `README.md` for the upstream source and license reference.

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

use crate::RuntimeTarget;

/// Extra information when generating the `metadata-hash`.
#[cfg(feature = "metadata-hash")]
pub(crate) struct MetadataExtraInfo {
    pub decimals: u8,
    pub token_symbol: String,
}

/// Returns the manifest dir from the `CARGO_MANIFEST_DIR` env.
fn get_manifest_dir() -> PathBuf {
    env::var("CARGO_MANIFEST_DIR")
        .expect("`CARGO_MANIFEST_DIR` is always set for `build.rs` files; qed")
        .into()
}

/// First step of the [`WasmBuilder`] to select the project to build.
pub struct WasmBuilderSelectProject {
    /// This parameter just exists to make it impossible to construct
    /// this type outside of this crate.
    _ignore: (),
}

impl WasmBuilderSelectProject {
    /// Use the current project as project for building the WASM binary.
    ///
    /// # Panics
    ///
    /// Panics if the `CARGO_MANIFEST_DIR` variable is not set. This variable
    /// is always set by `Cargo` in `build.rs` files.
    pub fn with_current_project(self) -> WasmBuilder {
        WasmBuilder {
            cargo_flags: Vec::new(),
            rust_flags: Vec::new(),
            file_name: None,
            project_cargo_toml: get_manifest_dir().join("Cargo.toml"),
            features_to_enable: Vec::new(),
            disable_runtime_version_section_check: false,
            export_heap_base: false,
            import_memory: false,
            #[cfg(feature = "metadata-hash")]
            enable_metadata_hash: None,
        }
    }

    /// Use the given `path` as project for building the WASM binary.
    ///
    /// Returns an error if the given `path` does not points to a `Cargo.toml`.
    pub fn with_project(self, path: impl Into<PathBuf>) -> Result<WasmBuilder, &'static str> {
        let path = path.into();

        if path.ends_with("Cargo.toml") && path.exists() {
            Ok(WasmBuilder {
                cargo_flags: Vec::new(),
                rust_flags: Vec::new(),
                file_name: None,
                project_cargo_toml: path,
                features_to_enable: Vec::new(),
                disable_runtime_version_section_check: false,
                export_heap_base: false,
                import_memory: false,
                #[cfg(feature = "metadata-hash")]
                enable_metadata_hash: None,
            })
        } else {
            Err("Project path must point to the `Cargo.toml` of the project")
        }
    }
}

/// The builder for building a wasm binary.
///
/// The builder itself is separated into multiple structs to make the setup type safe.
///
/// Building a wasm binary:
///
/// 1. Call [`WasmBuilder::new`] to create a new builder.
/// 2. Select the project to build using the methods of [`WasmBuilderSelectProject`].
/// 3. Set additional `RUST_FLAGS` or a different name for the file containing the WASM code using
///    methods of [`WasmBuilder`].
/// 4. Build the WASM binary using [`Self::build`].
pub struct WasmBuilder {
    /// Flags that should be appended to `cargo rustc` invocation.
    cargo_flags: Vec<String>,
    /// Flags that should be appended to `RUST_FLAGS` env variable.
    rust_flags: Vec<String>,
    /// The name of the file that is being generated in `OUT_DIR`.
    ///
    /// Defaults to `wasm_binary.rs`.
    file_name: Option<String>,
    /// The path to the `Cargo.toml` of the project that should be built
    /// for wasm.
    project_cargo_toml: PathBuf,
    /// Features that should be enabled when building the wasm binary.
    features_to_enable: Vec<String>,
    /// Should the builder not check that the `runtime_version` section exists in the wasm binary?
    disable_runtime_version_section_check: bool,

    /// Whether `__heap_base` should be exported (WASM-only).
    export_heap_base: bool,
    /// Whether `--import-memory` should be added to the link args (WASM-only).
    import_memory: bool,

    /// Whether to enable the metadata hash generation.
    #[cfg(feature = "metadata-hash")]
    enable_metadata_hash: Option<MetadataExtraInfo>,
}

impl WasmBuilder {
    /// Create a new instance of the builder.
    pub fn new() -> WasmBuilderSelectProject {
        WasmBuilderSelectProject { _ignore: () }
    }

    /// Build the WASM binary using the recommended default values.
    ///
    /// This is the same as calling:
    /// ```no_run
    /// substrate_wasm_builder::WasmBuilder::new()
    ///     .with_current_project()
    ///     .import_memory()
    ///     .export_heap_base()
    ///     .build();
    /// ```
    pub fn build_using_defaults() {
        WasmBuilder::new()
            .with_current_project()
            .import_memory()
            .export_heap_base()
            .build();
    }

    /// Init the wasm builder with the recommended default values.
    ///
    /// In contrast to [`Self::build_using_defaults`] it does not build the WASM binary directly.
    ///
    /// This is the same as calling:
    /// ```no_run
    /// substrate_wasm_builder::WasmBuilder::new()
    ///     .with_current_project()
    ///     .import_memory()
    ///     .export_heap_base();
    /// ```
    pub fn init_with_defaults() -> Self {
        WasmBuilder::new()
            .with_current_project()
            .import_memory()
            .export_heap_base()
    }

    /// Enable exporting `__heap_base` as global variable in the WASM binary.
    ///
    /// This adds `-C link-arg=--export=__heap_base` to `RUST_FLAGS`.
    pub fn export_heap_base(mut self) -> Self {
        self.export_heap_base = true;
        self
    }

    /// Set the name of the file that will be generated in `OUT_DIR`.
    ///
    /// This file needs to be included to get access to the build WASM binary.
    ///
    /// If this function is not called, `file_name` defaults to `wasm_binary.rs`
    pub fn set_file_name(mut self, file_name: impl Into<String>) -> Self {
        self.file_name = Some(file_name.into());
        self
    }

    /// Instruct the linker to import the memory into the WASM binary.
    ///
    /// This adds `-C link-arg=--import-memory` to `RUST_FLAGS`.
    pub fn import_memory(mut self) -> Self {
        self.import_memory = true;
        self
    }

    /// Append the given `flag` to `cargo rustc` invocation.
    ///
    /// `flag` is appended as is, so it needs to be a valid flag.
    pub fn append_to_cargo_flags(mut self, flag: impl Into<String>) -> Self {
        self.cargo_flags.push(flag.into());
        self
    }

    /// Append the given `flag` to `RUST_FLAGS`.
    ///
    /// `flag` is appended as is, so it needs to be a valid flag.
    pub fn append_to_rust_flags(mut self, flag: impl Into<String>) -> Self {
        self.rust_flags.push(flag.into());
        self
    }

    /// Enable the given feature when building the wasm binary.
    ///
    /// `feature` needs to be a valid feature that is defined in the project `Cargo.toml`.
    pub fn enable_feature(mut self, feature: impl Into<String>) -> Self {
        self.features_to_enable.push(feature.into());
        self
    }

    /// Enable generation of the metadata hash.
    ///
    /// This will compile the runtime once, fetch the metadata, build the metadata hash and
    /// then compile again with the env `RUNTIME_METADATA_HASH` set. For more information
    /// about the metadata hash see [RFC78](https://polkadot-fellows.github.io/RFCs/approved/0078-merkleized-metadata.html).
    ///
    /// - `token_symbol`: The symbol of the main native token of the chain.
    /// - `decimals`: The number of decimals of the main native token.
    #[cfg(feature = "metadata-hash")]
    pub fn enable_metadata_hash(mut self, token_symbol: impl Into<String>, decimals: u8) -> Self {
        self.enable_metadata_hash = Some(MetadataExtraInfo {
            token_symbol: token_symbol.into(),
            decimals,
        });

        self
    }

    /// Disable the check for the `runtime_version` wasm section.
    ///
    /// By default the `wasm-builder` will ensure that the `runtime_version` section will
    /// exists in the build wasm binary. This `runtime_version` section is used to get the
    /// `RuntimeVersion` without needing to call into the wasm binary. However, for some
    /// use cases (like tests) you may want to disable this check.
    pub fn disable_runtime_version_section_check(mut self) -> Self {
        self.disable_runtime_version_section_check = true;
        self
    }

    /// Build the WASM binary.
    pub fn build(mut self) {
        let target = crate::runtime_target();
        if target == RuntimeTarget::Wasm {
            if self.export_heap_base {
                self.rust_flags
                    .push("-C link-arg=--export=__heap_base".into());
            }

            if self.import_memory {
                self.rust_flags.push("-C link-arg=--import-memory".into());
            }
        }

        let out_dir = PathBuf::from(env::var("OUT_DIR").expect("`OUT_DIR` is set by cargo!"));
        let file_path = out_dir.join(
            self.file_name
                .clone()
                .unwrap_or_else(|| "wasm_binary.rs".into()),
        );

        if check_skip_build() {
            // If we skip the build, we still want to make sure to be called when an env variable
            // changes
            generate_rerun_if_changed_instructions();

            provide_dummy_wasm_binary_if_not_exist(&file_path);

            return;
        }

        build_project(
            target,
            file_path,
            self.project_cargo_toml,
            self.cargo_flags,
            self.rust_flags.join(" "),
            self.features_to_enable,
            self.file_name,
            !self.disable_runtime_version_section_check,
            #[cfg(feature = "metadata-hash")]
            self.enable_metadata_hash,
        );

        // As last step we need to generate our `rerun-if-changed` stuff. If a build fails, we don't
        // want to spam the output!
        generate_rerun_if_changed_instructions();
    }
}

/// Generate the name of the skip build environment variable for the current crate.
fn generate_crate_skip_build_env_name() -> String {
    format!(
        "SKIP_{}_WASM_BUILD",
        env::var("CARGO_PKG_NAME")
            .expect("Package name is set")
            .to_uppercase()
            .replace('-', "_"),
    )
}

/// Checks if the build of the WASM binary should be skipped.
fn check_skip_build() -> bool {
    env::var(crate::SKIP_BUILD_ENV).is_ok() ||
		env::var(generate_crate_skip_build_env_name()).is_ok() ||
		// If we are running in docs.rs, let's skip building.
		// https://docs.rs/about/builds#detecting-docsrs
		env::var("DOCS_RS").is_ok()
}

/// Provide a dummy WASM binary if there doesn't exist one.
fn provide_dummy_wasm_binary_if_not_exist(file_path: &Path) {
    if !file_path.exists() {
        crate::write_file_if_changed(
            file_path,
            "pub const WASM_BINARY: Option<&[u8]> = None;\
			 pub const WASM_BINARY_BLOATY: Option<&[u8]> = None;",
        );
    }
}

/// Generate the `rerun-if-changed` instructions for cargo to make sure that the WASM binary is
/// rebuilt when needed.
fn generate_rerun_if_changed_instructions() {
    // Make sure that the `build.rs` is called again if one of the following env variables changes.
    println!("cargo:rerun-if-env-changed={}", crate::SKIP_BUILD_ENV);
    println!("cargo:rerun-if-env-changed={}", crate::FORCE_WASM_BUILD_ENV);
    println!(
        "cargo:rerun-if-env-changed={}",
        generate_crate_skip_build_env_name()
    );
}

/// Build the currently built project as wasm binary.
///
/// The current project is determined by using the `CARGO_MANIFEST_DIR` environment variable.
///
/// `file_name` - The name + path of the file being generated. The file contains the
/// constant `WASM_BINARY`, which contains the built wasm binary.
///
/// `project_cargo_toml` - The path to the `Cargo.toml` of the project that should be built.
///
/// `default_rustflags` - Default `RUSTFLAGS` that will always be set for the build.
///
/// `features_to_enable` - Features that should be enabled for the project.
///
/// `wasm_binary_name` - The optional wasm binary name that is extended with
/// `.compact.compressed.wasm`. If `None`, the project name will be used.
///
/// `check_for_runtime_version_section` - Should the wasm binary be checked for the
/// `runtime_version` section?
fn build_project(
    target: RuntimeTarget,
    file_name: PathBuf,
    project_cargo_toml: PathBuf,
    default_cargo_flags: Vec<String>,
    default_rustflags: String,
    features_to_enable: Vec<String>,
    wasm_binary_name: Option<String>,
    check_for_runtime_version_section: bool,
    #[cfg(feature = "metadata-hash")] enable_metadata_hash: Option<MetadataExtraInfo>,
) {
    // Init jobserver as soon as possible
    crate::wasm_project::get_jobserver();
    let cargo_cmd = match crate::prerequisites::check(target) {
        Ok(cmd) => cmd,
        Err(err_msg) => {
            eprintln!("{err_msg}");
            process::exit(1);
        }
    };

    let (wasm_binary, bloaty) = crate::wasm_project::create_and_compile(
        target,
        &project_cargo_toml,
        &default_cargo_flags,
        &default_rustflags,
        cargo_cmd,
        features_to_enable,
        wasm_binary_name,
        check_for_runtime_version_section,
        #[cfg(feature = "metadata-hash")]
        enable_metadata_hash,
    );

    let (wasm_binary, wasm_binary_bloaty) = if let Some(wasm_binary) = wasm_binary {
        (
            wasm_binary.wasm_binary_path_escaped(),
            bloaty.bloaty_path_escaped(),
        )
    } else {
        (bloaty.bloaty_path_escaped(), bloaty.bloaty_path_escaped())
    };

    crate::write_file_if_changed(
        file_name,
        format!(
            r#"
				pub const WASM_BINARY: Option<&[u8]> = Some(include_bytes!("{wasm_binary}"));
				pub const WASM_BINARY_BLOATY: Option<&[u8]> = Some(include_bytes!("{wasm_binary_bloaty}"));
			"#,
            wasm_binary = wasm_binary,
            wasm_binary_bloaty = wasm_binary_bloaty,
        ),
    );
}
