imbibe-protos 0.0.1

#[path = "codegen/signer.rs"]
pub mod signer;

use std::{
	collections::BTreeMap,
	fs::{self, File},
	io::{self, BufWriter, Write},
	path::Path,
};

use anyhow::Context;
use prost_build::Config;
use prost_reflect::DescriptorPool;
use walkdir::WalkDir;

use crate::global::OUT_DIR;

const GENERATED_DO_NOT_EDIT: &str = "// This is @generated by build.rs. DO NOT EDIT directly.\n";

#[derive(Default, Debug)]
pub struct ModuleNode {
	submodules: BTreeMap<String, ModuleNode>,
	src_file_name: Option<String>,
}

pub fn compile_protos<P1, P2>(proto_dir: P1, gen_dir: P2) -> anyhow::Result<()>
where
	P1: AsRef<Path>,
	P2: AsRef<Path>,
{
	let protos: Vec<_> = WalkDir::new(proto_dir.as_ref())
		.into_iter()
		.filter_map(|proto| proto.ok())
		.filter(|e| e.path().extension().is_some_and(|ext| ext == "proto"))
		.map(|e| e.path().to_path_buf())
		.collect();

	if protos.is_empty() {
		println!(
			"cargo:warning=no .proto files found in '{}'",
			proto_dir.as_ref().display(),
		);

		// only generated to avoid compilation failure when no proto files exported
		let mut file = File::create(OUT_DIR.join("any_signer_extractor.rs")).map(BufWriter::new)?;
		writeln!(file, "{GENERATED_DO_NOT_EDIT}")?;

		return Ok(());
	};

	println!("running prost-build");

	let pool = file_descriptor_pool(&protos, &proto_dir, &gen_dir)?;
	let signer_msgs = signer::find_signer_msgs(&pool)?;
	generate_signer_extractors_macro_invocation(signer_msgs.keys())?;

	let mut config = Config::new();
	signer_msgs.into_iter().for_each(|(msg_name, fields)| {
		config.type_attribute(&msg_name, "#[derive(GetSigners)]");
		config.type_attribute(msg_name, format!("#[signer_fields({})]", fields.join(",")));
	});

	config
		.out_dir(gen_dir.as_ref())
		.enable_type_names()
		.protoc_executable(protoc_bin_vendored::protoc_bin_path()?)
		.compile_protos(&protos, &[proto_dir])
		.context("prost-build compilation failed")?;

	println!("info: prost-build finished successfully");

	Ok(())
}

pub fn generate_mod_rs<P>(gen_dir: P) -> anyhow::Result<()>
where
	P: AsRef<Path>,
{
	let mut root_module_node = ModuleNode::default();
	let mut found_rs_files_count = 0;

	let gen_dir = gen_dir.as_ref();

	let mod_rs_path = gen_dir.join("mod.rs");
	let mut mod_rs_file = BufWriter::new(File::create(&mod_rs_path)?);

	if !gen_dir.try_exists()? {
		println!(
			"cargo:warning=PROTO_GEN_DIR ({}) does not exist, generated mod.rs will be empty regarding protos",
			gen_dir.display(),
		);

		writeln!(
			mod_rs_file,
			"// no Protobuf Rust files found in PROTO_GEN_DIR ({}) during build",
			gen_dir.display(),
		)?;

		mod_rs_file.flush()?;

		return Ok(());
	}

	println!(
		"scanning PROTO_GEN_DIR ({}) for generated .rs files to build mod.rs...",
		gen_dir.display(),
	);

	for entry in fs::read_dir(gen_dir)? {
		let entry = entry?;
		let path = entry.path();
		if path.is_file() && path.extension().is_some_and(|ext| ext == "rs") {
			let original_file_name = path
				.file_name()
				.with_context(|| format!("failed to get file name for path: {}", path.display()))?
				.to_str()
				.with_context(|| format!("file name is not valid UTF-8: {:?}", path.as_os_str()))?;

			if original_file_name == "mod.rs" {
				continue;
			}
			found_rs_files_count += 1;

			// package parts from filename (e.g., "cosmos.bank.v1beta1.rs" -> ["cosmos", "bank", "v1beta1"])
			let package_string = original_file_name.trim_end_matches(".rs");
			let parts: Vec<String> = package_string.split('.').map(String::from).collect();

			if !parts.is_empty() {
				println!(
					"processing for mod.rs: file='{}', package_parts='{:?}'",
					original_file_name, parts,
				);
				insert_into_tree(&mut root_module_node, &parts, original_file_name.into());
			} else {
				println!(
					"cargo:warning=skipping file for mod.rs (no package parts derived): '{}'",
					original_file_name,
				);
			}
		}
	}

	writeln!(mod_rs_file, "{GENERATED_DO_NOT_EDIT}")?;

	found_rs_files_count.eq(&0).then(||{
		println!(
			"cargo:warning=No .rs files (other than potential mod.rs itself) found in {} to include in generated mod.rs. This might be unexpected if protos were supposed to be compiled.",
			gen_dir.display()
		);

		writeln!(
			mod_rs_file,
			"// No Protobuf Rust files were found in PROTO_GEN_DIR during build."
		)
	})
	.transpose()?
	.is_none()
	.then(|| {
		writeln!(
			mod_rs_file,
			"// It includes all Rust modules generated from .proto files, structured hierarchically using include!.\n"
		)?;
		write_tree_to_mod_rs_recursive(
			&mut mod_rs_file,
			&root_module_node.submodules,
			0,
			&mut vec![],
		)
	})
	.transpose()?;

	mod_rs_file.flush()?;

	println!(
		"generated nested mod.rs with include! strategy at {}",
		mod_rs_path.display()
	);

	Ok(())
}

fn file_descriptor_pool<P1, P2, P3>(
	protos: &[P1],
	proto_dir: P2,
	gen_dir: P3,
) -> anyhow::Result<DescriptorPool>
where
	P1: AsRef<Path>,
	P2: AsRef<Path>,
	P3: AsRef<Path>,
{
	let fds_path = OUT_DIR.join("file_descriptor_set.bin");
	Config::new()
		.out_dir(gen_dir.as_ref())
		.enable_type_names()
		.file_descriptor_set_path(&fds_path)
		.protoc_executable(protoc_bin_vendored::protoc_bin_path()?)
		.compile_protos(protos, &[proto_dir])
		.context("failed to generate file descriptor set")?;

	let descriptor_bytes = fs::read(fds_path)?;
	DescriptorPool::decode(descriptor_bytes.as_slice()).map_err(From::from)
}

fn generate_signer_extractors_macro_invocation<I>(signer_msgs: I) -> anyhow::Result<()>
where
	I: Iterator<Item: AsRef<str>>,
{
	let mut file = BufWriter::new(File::create(OUT_DIR.join("any_signer_extractor.rs"))?);

	writeln!(file, "{GENERATED_DO_NOT_EDIT}")?;

	writeln!(file, "generate_signer_extractors!(")?;
	for msg_name in signer_msgs {
		let type_url = format!("/{}", msg_name.as_ref());
		let rust_path = msg_name.as_ref().replace('.', "::");
		writeln!(file, "\t(\"{type_url}\", {rust_path}),")?;
	}

	writeln!(file, ");")?;

	Ok(())
}

fn insert_into_tree(root: &mut ModuleNode, parts: &[String], original_file_name: String) {
	let mut current_node = root;
	for part in parts {
		current_node = current_node.submodules.entry(part.clone()).or_default();
	}
	current_node.src_file_name = Some(original_file_name.to_string());
}

fn write_tree_to_mod_rs_recursive(
	writer: &mut BufWriter<File>,
	module_map: &BTreeMap<String, ModuleNode>,
	indent_level: usize,
	current_path_segments: &mut Vec<String>,
) -> io::Result<()> {
	let indent = "\t".repeat(indent_level);

	for (mod_segment_name, node_data) in module_map {
		current_path_segments.push(mod_segment_name.clone());
		let rust_mod_name = mod_segment_name.to_lowercase();

		writeln!(writer, "{}pub mod {} {{", indent, rust_mod_name)?;

		match node_data.src_file_name.as_ref() {
			Some(src_file_name) => {
				let indent = format!("{indent}\t");
				writeln!(writer, "{}#[allow(unused_imports)]", indent)?;
				writeln!(writer, "{}use crate::GetSigners;\n", indent)?;
				writeln!(writer, "{}include!(\"{}\");", indent, src_file_name)?;
			},
			None => {
				println!(
					"mod.rs entry for module path '{}': pub mod {} {{ ... }} (namespace only)",
					current_path_segments.join("::"),
					rust_mod_name,
				);
			},
		}

		// Recursively define submodules, if any, within the current module block.
		// This handles cases where a package like "a.b" might have generated "a.b.rs"
		// AND there are further sub-packages like "a.b.c" which generates "a.b.c.rs".
		// The "a.b.c.rs" content would be included inside the "c" module, within "b".
		if !node_data.submodules.is_empty() {
			write_tree_to_mod_rs_recursive(
				writer,
				&node_data.submodules,
				indent_level + 1,
				current_path_segments,
			)?;
		}

		writeln!(writer, "{}}}\n", indent)?; // Close the current module block
		current_path_segments.pop();
	}

	Ok(())
}