camber-build 0.1.0

use std::io;
use std::path::Path;

/// Compile `.proto` files and generate async service traits.
///
/// Runs `tonic-build` to produce standard async gRPC server code, then
/// generates an additional async trait and bridge struct per service.
/// The trait has async methods. The bridge delegates directly via `.await`.
///
/// Call this from your crate's `build.rs`.
pub fn compile_protos(
    protos: &[impl AsRef<Path>],
    includes: &[impl AsRef<Path>],
) -> io::Result<()> {
    configure().compile_protos(protos, includes)
}

/// Return a builder for customizing proto compilation.
pub fn configure() -> Builder {
    Builder {
        file_descriptor_set_path: None,
    }
}

pub struct Builder {
    file_descriptor_set_path: Option<std::path::PathBuf>,
}

impl Builder {
    /// Generate a file containing the encoded `FileDescriptorSet`.
    /// Required for gRPC reflection.
    pub fn file_descriptor_set_path(mut self, path: impl AsRef<Path>) -> Self {
        self.file_descriptor_set_path = Some(path.as_ref().to_path_buf());
        self
    }

    /// Compile the `.proto` files and generate code.
    pub fn compile_protos(
        self,
        protos: &[impl AsRef<Path>],
        includes: &[impl AsRef<Path>],
    ) -> io::Result<()> {
        let tonic = tonic_build::configure()
            .build_client(true)
            .build_server(true)
            .build_transport(false);

        let mut config = prost_build::Config::new();
        if let Some(path) = &self.file_descriptor_set_path {
            config.file_descriptor_set_path(path);
        }
        config.service_generator(Box::new(ServiceGenerator::new(tonic)));
        config.compile_protos(protos, includes)
    }
}

/// Generates both standard tonic async code and camber async wrappers.
struct ServiceGenerator {
    tonic_gen: Box<dyn prost_build::ServiceGenerator>,
    service_code: String,
}

impl ServiceGenerator {
    fn new(tonic_builder: tonic_build::Builder) -> Self {
        Self {
            tonic_gen: tonic_builder.service_generator(),
            service_code: String::new(),
        }
    }
}

impl prost_build::ServiceGenerator for ServiceGenerator {
    fn generate(&mut self, service: prost_build::Service, buf: &mut String) {
        self.tonic_gen.generate(service.clone(), buf);
        generate_service_wrapper(&service, &mut self.service_code);
    }

    fn finalize(&mut self, buf: &mut String) {
        self.tonic_gen.finalize(buf);

        if !self.service_code.is_empty() {
            buf.push_str(&self.service_code);
            self.service_code.clear();
        }
    }

    fn finalize_package(&mut self, package: &str, buf: &mut String) {
        self.tonic_gen.finalize_package(package, buf);
    }
}

fn snake_case(name: &str) -> String {
    let mut s = String::new();
    let mut it = name.chars().peekable();
    while let Some(c) = it.next() {
        s.push(c.to_ascii_lowercase());
        if it.peek().is_some_and(|next| next.is_uppercase()) {
            s.push('_');
        }
    }
    s
}

/// Build the function signature for a unary RPC method.
/// Returns `None` for streaming methods (not yet supported).
fn unary_signature(method: &prost_build::Method) -> Option<String> {
    match (method.client_streaming, method.server_streaming) {
        (false, false) => {
            let name = &method.name;
            let input = &method.input_type;
            let output = &method.output_type;
            Some(format!(
                "async fn {name}(&self, request: tonic::Request<{input}>) -> \
                 std::result::Result<tonic::Response<{output}>, tonic::Status>"
            ))
        }
        _ => None,
    }
}

fn generate_service_wrapper(service: &prost_build::Service, buf: &mut String) {
    let trait_name = &service.name;
    let base = snake_case(trait_name);
    let mod_name = format!("{base}_service");
    let server_type = format!("{}Server", trait_name);
    let server_mod = format!("{base}_server");

    buf.push_str(&format!(
        "/// Async wrappers for the {trait_name} service.\n"
    ));
    buf.push_str(&format!("pub mod {mod_name} {{\n"));
    buf.push_str("    use super::*;\n\n");

    // Generate async trait
    buf.push_str(&format!(
        "    /// Async `{trait_name}` gRPC service trait.\n"
    ));
    buf.push_str("    /// Implement this instead of the tonic-generated version.\n");
    buf.push_str("    #[tonic::async_trait]\n");
    buf.push_str(&format!(
        "    pub trait {trait_name}: Send + Sync + 'static {{\n"
    ));

    for method in &service.methods {
        match unary_signature(method) {
            Some(sig) => {
                buf.push_str(&format!("        {sig};\n"));
            }
            None => {
                let method_name = &method.name;
                buf.push_str(&format!(
                    "        compile_error!(\"camber-build: streaming RPCs are not yet supported (method `{method_name}` in service `{trait_name}`)\");\n"
                ));
            }
        }
    }
    buf.push_str("    }\n\n");

    // Generate bridge struct
    buf.push_str(&format!(
        "    /// Bridges an async `{trait_name}` impl to the tonic service trait.\n"
    ));
    buf.push_str(&format!(
        "    pub struct Bridge<T: {trait_name}>(pub T);\n\n"
    ));

    // Implement the tonic async trait via the bridge
    buf.push_str("    #[tonic::async_trait]\n");
    buf.push_str(&format!(
        "    impl<T: {trait_name}> {server_mod}::{trait_name} for Bridge<T> {{\n"
    ));

    for method in &service.methods {
        if let Some(sig) = unary_signature(method) {
            let method_name = &method.name;
            buf.push_str(&format!("        {sig} {{\n"));
            buf.push_str(&format!(
                "            self.0.{method_name}(request).await\n"
            ));
            buf.push_str("        }\n");
        }
    }
    buf.push_str("    }\n\n");

    // Helper function to create a tonic server from an async impl
    buf.push_str(&format!(
        "    /// Create a tonic `{server_type}` from an async service implementation.\n"
    ));
    buf.push_str(&format!(
        "    pub fn serve<T: {trait_name}>(service: T) -> {server_mod}::{server_type}<Bridge<T>> {{\n"
    ));
    buf.push_str(&format!(
        "        {server_mod}::{server_type}::new(Bridge(service))\n"
    ));
    buf.push_str("    }\n");

    buf.push_str("}\n");
}