//! Builder logic for creating PVM code blobs for execution on the JAM PVM instances (service code
//! and authorizer code).

#![allow(clippy::unwrap_used)]

use std::{fs, path::PathBuf, process::Command};

pub enum BlobType {
	Service,
	Authorizer,
}

impl BlobType {
	pub fn dispatch_table(&self) -> Vec<Vec<u8>> {
		match self {
			Self::Service =>
				vec![b"refine_ext".into(), b"accumulate_ext".into(), b"on_transfer_ext".into()],
			Self::Authorizer => vec![b"is_authorized_ext".into()],
		}
	}
}

/// Build the given project for `riscv`.
///
/// A file `file_name` will be created in `OUT_DIR` that contains a constant named `RISCV_BLOB:
/// Option<&[u8]>` that represents the build riscv binary.
///
/// When setting the `SKIP_RUNTIME_BUILD` env variable at build time, the `RISCV_BLOB` will be set
/// to `None`.
pub fn build_jam_pvm_blob(path_name: &str, crate_name: &str, file_name: &str, blob_type: BlobType) {
	let out_dir: PathBuf = std::env::var("OUT_DIR").expect("No OUT_DIR").into();

	if std::env::var("SKIP_RUNTIME_BUILD").is_ok() {
		fs::write(out_dir.join(file_name), "pub const RISCV_BLOB: Option<&[u8]> = None;")
			.expect("Writes the output file.");
		return;
	}

	println!("cargo:rerun-if-env-changed=SKIP_RUNTIME_BUILD");

	let mut root_dir: PathBuf = std::env::var("CARGO_MANIFEST_DIR").unwrap().into();
	root_dir.pop();
	root_dir.pop();
	let service_dir = root_dir.join("services").join(path_name);
	let jam_common_dir = root_dir.join("crates").join("jam-pvm-common");

	let is_64_bit = cfg!(feature = "use-64-bit-pvm");
	let (target_name, target_json_path) = if is_64_bit {
		("riscv64emac-unknown-none-polkavm", polkavm_linker::target_json_64_path().unwrap())
	} else {
		("riscv32emac-unknown-none-polkavm", polkavm_linker::target_json_32_path().unwrap())
	};

	let build_res = Command::new("cargo")
		.current_dir(&service_dir)
		.env_clear()
		.env("PATH", std::env::var("PATH").unwrap())
		.env("RUSTFLAGS", "-C panic=abort")
		.env("CARGO_TARGET_DIR", &out_dir)
		.args(["build", "-Z", "build-std=core,alloc", "--release", "--target"])
		.arg(target_json_path)
		.output()
		.expect("Failed to execute rustup process");

	if !build_res.status.success() {
		let stderr = String::from_utf8_lossy(&build_res.stderr);
		eprintln!("{}", stderr);
		std::process::exit(1);
	}

	// Post processing
	let mut config = polkavm_linker::Config::default();
	config.set_strip(true);
	config.set_dispatch_table(blob_type.dispatch_table());
	let input_path = &out_dir.join(target_name).join("release").join(crate_name);
	let orig =
		fs::read(input_path).unwrap_or_else(|e| panic!("Failed to read {:?} :{:?}", input_path, e));
	let linked = polkavm_linker::program_from_elf(config, orig.as_ref())
		.expect("Failed to link polkavm program:");

	// Write out a full `.polkavm` blob for debugging/inspection.
	let output_path_polkavm = &out_dir.join(format!("{}.polkavm", crate_name));
	fs::write(output_path_polkavm, &linked).expect("Error writing resulting binary");

	let parts = polkavm_linker::ProgramParts::from_bytes(linked.into())
		.expect("failed to deserialize linked PolkaVM program");

	let mut ro_data = parts.ro_data.to_vec();
	ro_data.resize(parts.ro_data_size as usize, 0);

	let blob_jam = jam_types::ProgramBlob {
		ro_data: ro_data.into(),
		rw_data: (&parts.rw_data[..]).into(),
		code_blob: (&parts.code_and_jump_table[..]).into(),
		rw_data_padding: (parts.rw_data_size as usize - parts.rw_data.len()) as u32,
		stack_size: parts.stack_size,
	};

	let output_path_jam = &out_dir.join(format!("{}.jam", crate_name));
	fs::write(output_path_jam, blob_jam.to_vec().expect("error serializing the .jam blob"))
		.expect("error writing the .jam blob");

	fs::write(
		out_dir.join(file_name),
		format!(
			"pub const RISCV_BLOB: Option<&[u8]> = Some(include_bytes!(\"{}\"));",
			output_path_jam.display()
		),
	)
	.expect("Writes the output file.");

	println!("cargo:rerun-if-changed={}", service_dir.to_str().unwrap());
	println!("cargo:rerun-if-changed={}", jam_common_dir.to_str().unwrap());
}

/// Build the service crate in the parent path for the RISCV target, convert to PVM code and finish
/// by creating a `.pvm` blob file for inclusion into the `lib.rs` of the present crate using the
/// `pvm_binary` macro.
pub fn build_main() {
	let mut crate_dir: PathBuf = std::env::var("CARGO_MANIFEST_DIR").unwrap().into();
	crate_dir.pop();
	let this_crate_name = std::env::var("CARGO_PKG_NAME").expect("No PACKAGE_NAME");
	let crate_name = this_crate_name
		.strip_suffix("-bin")
		.expect("This build script is only intended for binary crates ending in `-bin`");
	if crate_name.ends_with("-service") {
		build_service_main(crate_dir, crate_name);
	} else if crate_name.ends_with("-authorizer") {
		build_authorizer_main(crate_dir, crate_name);
	} else {
		panic!("This function is only for building crates ending in `-service` or `-authorizer`");
	}
}

/// Build the service crate in the parent path for the RISCV target, convert to PVM code and finish
/// by creating a `.pvm` blob file for inclusion into the `lib.rs` of the present crate using the
/// `pvm_binary` macro.
pub fn build_service_main(crate_dir: PathBuf, crate_name: &str) {
	let this_crate_name = std::env::var("CARGO_PKG_NAME").expect("No PACKAGE_NAME");
	let out_dir: PathBuf = std::env::var("OUT_DIR").expect("No OUT_DIR").into();
	let output_file = out_dir.join(format!("{}.jam", &this_crate_name));

	build_pvm_blob(crate_dir, crate_name, BlobType::Service, &output_file);
}

/// Build the authorizer crate in the parent path for the RISCV target, convert to PVM code and finish
/// by creating a `.pvm` blob file for inclusion into the `lib.rs` of the present crate using the
/// `pvm_binary` macro.
pub fn build_authorizer_main(crate_dir: PathBuf, crate_name: &str) {
	let this_crate_name = std::env::var("CARGO_PKG_NAME").expect("No PACKAGE_NAME");
	let out_dir: PathBuf = std::env::var("OUT_DIR").expect("No OUT_DIR").into();
	let output_file = out_dir.join(format!("{}.jam", &this_crate_name));

	build_pvm_blob(crate_dir, crate_name, BlobType::Authorizer, &output_file);
}

/// Build the PVM crate in `crate_dir` called `crate_name` for the RISCV target, convert to PVM
/// code and finish by creating a `.pvm` blob file for inclusion into the `lib.rs` of the present
/// crate (called `this_crate_name`) using the `pvm_binary` macro. The crate is either a service
/// or an authorizer, defined by `blob_type`.
pub fn build_pvm_blob(
	crate_dir: PathBuf,
	crate_name: &str,
	blob_type: BlobType,
	output_file: &PathBuf,
) {
	let out_dir: PathBuf = std::env::var("OUT_DIR").expect("No OUT_DIR").into();

	let is_64_bit = cfg!(feature = "use-64-bit-pvm");
	let (target_name, target_json_path) = if is_64_bit {
		("riscv64emac-unknown-none-polkavm", polkavm_linker::target_json_64_path().unwrap())
	} else {
		("riscv32emac-unknown-none-polkavm", polkavm_linker::target_json_32_path().unwrap())
	};

	let build_res = Command::new("cargo")
		.current_dir(&crate_dir)
		.env_clear()
		.env("PATH", std::env::var("PATH").unwrap())
		.env("RUSTFLAGS", "-C panic=abort")
		.env("CARGO_TARGET_DIR", &out_dir)
		.args(["build", "-Z", "build-std=core,alloc", "--release", "--target"])
		.arg(target_json_path)
		.output()
		.expect("Failed to execute rustup process");

	if !build_res.status.success() {
		let stderr = String::from_utf8_lossy(&build_res.stderr);
		eprintln!("{}", stderr);
		std::process::exit(1);
	}

	// Post processing
	let mut config = polkavm_linker::Config::default();
	config.set_strip(true);
	config.set_dispatch_table(blob_type.dispatch_table());
	let input_path = &out_dir.join(target_name).join("release").join(crate_name);
	let orig =
		fs::read(input_path).unwrap_or_else(|e| panic!("Failed to read {:?} :{:?}", input_path, e));
	let linked = polkavm_linker::program_from_elf(config, orig.as_ref())
		.expect("Failed to link polkavm program:");

	// Write out a full `.polkavm` blob for debugging/inspection.
	let output_path_polkavm = &out_dir.join(format!("{}.polkavm", &crate_name));
	fs::write(output_path_polkavm, &linked).expect("Error writing resulting binary");

	let parts = polkavm_linker::ProgramParts::from_bytes(linked.into())
		.expect("failed to deserialize linked PolkaVM program");

	let mut ro_data = parts.ro_data.to_vec();
	ro_data.resize(parts.ro_data_size as usize, 0);

	let blob_jam = jam_types::ProgramBlob {
		ro_data: ro_data.into(),
		rw_data: (&parts.rw_data[..]).into(),
		code_blob: (&parts.code_and_jump_table[..]).into(),
		rw_data_padding: (parts.rw_data_size as usize - parts.rw_data.len()) as u32,
		stack_size: parts.stack_size,
	};

	fs::write(output_file, blob_jam.to_vec().expect("error serializing the .jam blob"))
		.expect("error writing the .jam blob");

	println!("cargo:rerun-if-changed={}", crate_dir.to_str().unwrap());
}

#[macro_export]
macro_rules! pvm_binary {
	() => {
		include_bytes!(concat!(env!("OUT_DIR"), "/", env!("CARGO_PKG_NAME"), ".jam"));
	};
}
