// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use move_deps::{
    move_binary_format::{compatibility::Compatibility, normalized::Module, CompiledModule},
    move_command_line_common::files::{
        extension_equals, find_filenames, MOVE_COMPILED_EXTENSION, MOVE_ERROR_DESC_EXTENSION,
    },
    move_compiler::compiled_unit::{CompiledUnit, NamedCompiledModule},
    move_core_types::language_storage::ModuleId,
    move_package::{BuildConfig, ModelConfig},
};
use std::{
    collections::BTreeMap,
    fs::{create_dir_all, remove_dir_all, File},
    io::Read,
    path::{Path, PathBuf},
};
use structopt::*;

pub const CORE_FRAMEWORK_RELEASE_SUFFIX: &str = "framework";
pub const TOKEN_RELEASE_SUFFIX: &str = "token";

/// Options to configure the generation of a release.
#[derive(Debug, StructOpt, Clone)]
#[structopt(
    name = "Aptos Frameworks",
    about = "Release CLI for Aptos frameworks",
    author = "Aptos",
    rename_all = "kebab-case"
)]
pub struct ReleaseOptions {
    #[structopt(long = "no-check-linking-layout-compatibility")]
    pub no_check_layout_compatibility: bool,
    #[structopt(long = "no-build-docs")]
    pub no_build_docs: bool,
    #[structopt(long = "with-diagram")]
    pub with_diagram: bool,
    #[structopt(long = "no-script-builder")]
    pub no_script_builder: bool,
    #[structopt(long = "no-script-abi")]
    pub no_script_abis: bool,
    #[structopt(long = "no-errmap")]
    pub no_errmap: bool,
    #[structopt(
        long = "package",
        default_value = "aptos-framework",
        parse(from_os_str)
    )]
    pub package: PathBuf,
    #[structopt(long = "output", default_value = "current", parse(from_os_str))]
    pub output: PathBuf,
}

impl Default for ReleaseOptions {
    fn default() -> Self {
        Self {
            no_build_docs: true,
            package: PathBuf::from("aptos-framework"),
            no_check_layout_compatibility: false,
            with_diagram: false,
            no_script_abis: true,
            no_script_builder: true,
            no_errmap: true,
            output: PathBuf::from("current"),
        }
    }
}

impl ReleaseOptions {
    pub fn create_release(&self) {
        let output_path = self
            .package
            .join("releases")
            .join("artifacts")
            .join(&self.output);

        let mut old_module_apis = None;
        if !self.no_check_layout_compatibility {
            old_module_apis = extract_old_apis(&output_path);
        }

        if output_path.exists() {
            std::fs::remove_dir_all(&output_path).unwrap();
        }
        std::fs::create_dir_all(output_path.parent().unwrap()).unwrap();

        let build_config = move_deps::move_package::BuildConfig {
            generate_docs: !self.no_build_docs,
            generate_abis: !self.no_script_abis,
            install_dir: Some(output_path.clone()),
            ..Default::default()
        };

        let package_path = Path::new(std::env!("CARGO_MANIFEST_DIR")).join(&self.package);

        let compiled_package = build_config
            .clone()
            .compile_package(&package_path, &mut std::io::stdout())
            .unwrap();

        if !self.no_check_layout_compatibility {
            println!("Checking layout compatibility");
            if let Some(old_module_apis) = old_module_apis {
                let new_modules =
                    compiled_package
                        .all_compiled_units()
                        .filter_map(|unit| match unit {
                            CompiledUnit::Module(NamedCompiledModule { module, .. }) => {
                                Some(module)
                            }
                            CompiledUnit::Script(_) => None,
                        });
                check_api_compatibility(&old_module_apis, new_modules);
            }
        }

        if !self.no_errmap {
            println!("Generating error map");
            generate_error_map(&package_path, &output_path, build_config)
        }

        if !self.no_script_builder {
            println!("Generating script builders");
            let abi_paths: Vec<&Path> = vec![&output_path];
            generate_script_builder(&output_path.join("aptos_sdk_builder.rs"), &abi_paths[..])
        }
    }
}

fn recreate_dir(dir_path: impl AsRef<Path>) {
    let dir_path = dir_path.as_ref();
    remove_dir_all(&dir_path).unwrap_or(());
    create_dir_all(&dir_path).unwrap();
}

fn generate_error_map(package_path: &Path, output_path: &Path, build_config: BuildConfig) {
    let mut errmap_path = output_path
        .join("error_description")
        .join("error_description");
    errmap_path.set_extension(MOVE_ERROR_DESC_EXTENSION);

    recreate_dir(&errmap_path.parent().unwrap());

    let errmap_options = move_deps::move_errmapgen::ErrmapOptions {
        output_file: errmap_path.to_string_lossy().to_string(),
        ..Default::default()
    };

    let model = build_config
        .move_model_for_package(
            package_path,
            ModelConfig {
                target_filter: None,
                all_files_as_targets: true,
            },
        )
        .unwrap();

    let mut emapgen = move_deps::move_errmapgen::ErrmapGen::new(&model, &errmap_options);
    emapgen.gen();
    emapgen.save_result();
}

fn generate_script_builder(output_path: impl AsRef<Path>, abi_paths: &[&Path]) {
    let output_path = output_path.as_ref();

    let abis: Vec<_> = abi_paths
        .iter()
        .flat_map(|path| {
            aptos_sdk_builder::read_abis(&[path])
                .unwrap_or_else(|_| panic!("Failed to read ABIs at {}", path.to_string_lossy()))
        })
        .collect();

    {
        let mut file = std::fs::File::create(output_path)
            .expect("Failed to open file for Rust script build generation");
        aptos_sdk_builder::rust::output(&mut file, &abis, /* local types */ true)
            .expect("Failed to generate Rust builders");
    }

    std::process::Command::new("rustfmt")
        .arg("--config")
        .arg("imports_granularity=crate")
        .arg(output_path)
        .status()
        .expect("Failed to run rustfmt on generated code");
}

fn extract_old_apis(package_path: impl AsRef<Path>) -> Option<BTreeMap<ModuleId, Module>> {
    let modules_path = package_path.as_ref();

    if !modules_path.is_dir() {
        eprintln!(
            "Warning: failed to extract old module APIs -- path \"{}\" is not a directory",
            modules_path.to_string_lossy()
        );
        return None;
    }
    let mut old_module_apis = BTreeMap::new();
    let files = find_filenames(&[modules_path], |p| {
        extension_equals(p, MOVE_COMPILED_EXTENSION)
    })
    .unwrap();
    for f in files {
        let mut bytes = Vec::new();
        File::open(f)
            .expect("Failed to open module bytecode file")
            .read_to_end(&mut bytes)
            .expect("Failed to read module bytecode file");
        let m = CompiledModule::deserialize(&bytes).expect("Failed to deserialize module bytecode");
        old_module_apis.insert(m.self_id(), Module::new(&m));
    }
    Some(old_module_apis)
}

fn check_api_compatibility<'a, I>(old: &BTreeMap<ModuleId, Module>, new: I)
where
    I: Iterator<Item = &'a CompiledModule>,
{
    let mut is_linking_layout_compatible = true;
    for module in new {
        // extract new linking/layout API and check compatibility with old
        let new_module_id = module.self_id();
        if let Some(old_api) = old.get(&new_module_id) {
            let new_api = Module::new(module);
            let compatibility = Compatibility::check(old_api, &new_api);
            if is_linking_layout_compatible && !compatibility.is_fully_compatible() {
                println!("Found linking/layout-incompatible change:");
                is_linking_layout_compatible = false
            }
            if !compatibility.struct_and_function_linking {
                eprintln!("Linking API for structs/functions of module {} has changed. Need to redeploy all dependent modules.", new_module_id.name())
            }
            if !compatibility.struct_layout {
                eprintln!("Layout API for structs of module {} has changed. Need to do a data migration of published structs", new_module_id.name())
            }
        }
    }
}