tron-macros 0.1.10

use heck::ToSnakeCase;
use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{ItemImpl, parse_macro_input};

use std::path::PathBuf;

/// Trim trailing punctuation/whitespace, remove trailing ';' if present
fn trim_token(s: &str) -> String {
    s.trim().trim_end_matches(';').trim().to_string()
}

#[proc_macro_attribute]
pub fn auto_impl_bridge(attr: TokenStream, item: TokenStream) -> TokenStream {
    // Parse input impl block
    let input = parse_macro_input!(item as ItemImpl);
    let self_ty = &input.self_ty;

    // Ensure trait is present (impl Trait for Type)
    let trait_path = match input.trait_.as_ref() {
        Some((_, path, _)) => path.clone(),
        None => {
            return quote! {
                compile_error!("#[auto_impl_bridge] must be applied to `impl <Trait> for <Type> { }`");
            }
            .into();
        }
    };

    // parse attribute literal: #[auto_impl_bridge("../proto/bridge.proto")]
    let proto_lit = match syn::parse::<syn::LitStr>(attr) {
        Ok(lit) => lit.value(),
        Err(_) => {
            return quote! {
                compile_error!("expected a string literal path to proto, e.g. #[auto_impl_bridge(\"../proto/bridge.proto\")]");
            }
            .into();
        }
    };

    // Resolve path relative to crate root where macro is invoked
    let manifest = std::env::var("CARGO_MANIFEST_DIR").unwrap_or_else(|_| ".".to_string());
    let proto_path = PathBuf::from(&manifest).join(&proto_lit);

    let proto_src = match std::fs::read_to_string(&proto_path) {
        Ok(s) => s,
        Err(e) => {
            let msg = format!(
                "auto_impl_bridge: failed to read proto file at {}: {}. \n                 Provide correct path relative to crate root.",
                proto_path.display(),
                e
            );
            return quote! { compile_error!(#msg); }.into();
        }
    };

    let mut rpcs = Vec::new();
    let mut in_service = false;
    for line in proto_src.lines() {
        let line = line.trim();
        if line.starts_with("service Bridge") {
            in_service = true;
            continue;
        }
        if in_service && line.starts_with("}") {
            break;
        }
        if in_service && line.starts_with("rpc") {
            let after_rpc = line["rpc ".len()..].trim();
            let name_end = after_rpc.find('(').unwrap_or(after_rpc.len());
            let proto_name = after_rpc[..name_end].trim().to_string();
            if proto_name.is_empty() {
                continue;
            }

            let rest = &after_rpc[name_end..];
            let req_type = rest
                .split_once('(')
                .and_then(|(_, after)| after.split_once(')'))
                .map(|(inside, _)| inside.trim().to_string())
                .unwrap_or_default();

            let (req_stream, req_type) = if req_type.starts_with("stream ") {
                (true, trim_token(&req_type["stream ".len()..]))
            } else {
                (false, trim_token(&req_type))
            };

            let resp_section = rest
                .split_once("returns")
                .map(|(_, after)| after)
                .unwrap_or("");
            let resp_type = resp_section
                .split_once('(')
                .and_then(|(_, after)| after.split_once(')'))
                .map(|(inside, _)| inside.trim().to_string())
                .unwrap_or_default();

            let (resp_stream, resp_type) = if resp_type.starts_with("stream ") {
                (true, trim_token(&resp_type["stream ".len()..]))
            } else {
                (false, trim_token(&resp_type))
            };

            rpcs.push((proto_name, req_stream, req_type, resp_stream, resp_type));
        }
    }

    // Generate tokens for methods and associated types
    let mut assoc_types = Vec::new();
    let mut method_tokens = Vec::new();

    for (proto_name, req_stream, req_type, resp_stream, resp_type) in rpcs {
        // method identifier in snake_case
        let method_snake = proto_name.to_snake_case();
        let method_ident = format_ident!("{}", method_snake);
        let handler_ident = format_ident!("handle_{}", method_snake);

        // request/response type paths generated by tonic are typically in crate::bridge::<Type>
        // parse them as syn::Type
        let req_ty_str = format!("tron_protos::{}", req_type);
        let resp_ty_str = format!("tron_protos::{}", resp_type);

        let req_ty: syn::Type = match syn::parse_str(&req_ty_str) {
            Ok(t) => t,
            Err(_) => {
                let msg = format!(
                    "auto_impl_bridge: failed to parse request type `{}`",
                    req_ty_str
                );
                return quote! { compile_error!(#msg); }.into();
            }
        };

        let resp_ty: syn::Type = match syn::parse_str(&resp_ty_str) {
            Ok(t) => t,
            Err(_) => {
                let msg = format!(
                    "auto_impl_bridge: failed to parse response type `{}`",
                    resp_ty_str
                );
                return quote! { compile_error!(#msg); }.into();
            }
        };

        // If response is streaming, add an associated type alias:
        // type <ProtoName>Stream = crate::<ProtoName>Stream;
        if resp_stream {
            let assoc_ident = format_ident!("{}Stream", proto_name);
            let assoc_target = format_ident!("{}Stream", proto_name);
            assoc_types.push(quote! {
                type #assoc_ident = crate::#assoc_target;
            });
        }

        // Build method signature based on streaming flags
        let method_quote = match (req_stream, resp_stream) {
            (false, false) => {
                // unary -> unary
                quote! {
                    #[cfg_attr(any(debug_assertions, test), tracing::instrument(skip(self), fields(request = ?request.get_ref())))]
                    async fn #method_ident(
                        &self,
                        request: tonic::Request<#req_ty>,
                    ) -> Result<tonic::Response<#resp_ty>, tonic::Status> {
                        self.#handler_ident(request).await
                    }
                }
            }
            (true, false) => {
                // client streaming -> unary response
                quote! {
                    #[cfg_attr(any(debug_assertions, test), tracing::instrument(skip(self), fields(request = ?request.get_ref())))]
                    async fn #method_ident(
                        &self,
                        request: tonic::Request<tonic::Streaming<#req_ty>>,
                    ) -> Result<tonic::Response<#resp_ty>, tonic::Status> {
                        self.#handler_ident(request).await
                    }
                }
            }
            (false, true) => {
                // unary -> server streaming. return type uses associated type
                let assoc_ident = format_ident!("{}Stream", proto_name);
                quote! {
                    #[cfg_attr(any(debug_assertions, test), tracing::instrument(skip(self), fields(request = ?request.get_ref())))]
                    async fn #method_ident(
                        &self,
                        request: tonic::Request<#req_ty>,
                    ) -> Result<tonic::Response<Self::#assoc_ident>, tonic::Status> {
                        self.#handler_ident(request).await
                    }
                }
            }
            (true, true) => {
                // bidi streaming -> req is Streaming<T>, resp is associated stream type
                let assoc_ident = format_ident!("{}Stream", proto_name);
                quote! {
                    #[cfg_attr(any(debug_assertions, test), tracing::instrument(skip(self), fields(request = ?request.get_ref())))]
                    async fn #method_ident(
                        &self,
                        request: tonic::Request<tonic::Streaming<#req_ty>>,
                    ) -> Result<tonic::Response<Self::#assoc_ident>, tonic::Status> {
                        self.#handler_ident(request).await
                    }
                }
            }
        };

        method_tokens.push(method_quote);
    }

    // Emit the impl
    let expanded = quote! {
        #[tonic::async_trait]
        impl #trait_path for #self_ty {
            #(#assoc_types)*
            #(#method_tokens)*
        }
    };

    expanded.into()
}