flatbuffers-tonic-build 0.1.1

use std::path::Path;

use flatbuffers_util::reflect::{GeneratorContext, MessageType};
use heck::ToSnakeCase;
use proc_macro2::TokenStream;
use quote::quote;

pub(crate) fn compile_flatbuffers_tonic_file_only(
    path: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    let schema = flatbuffers_util::reflect::compile_reflection_schema(path);
    let schema_ref = schema.get_ref();
    let gen_ctx = flatbuffers_util::reflect::GeneratorContext::parse_from_schema(&schema_ref);
    let types = gen_ctx.collect_in_out_types();
    let services = gen_ctx.get_services();

    println!("Processing service: {services:?}");
    // println!("Processing types: {types:?}");

    // generate types alias file
    let out_dir = std::env::var("OUT_DIR").unwrap();
    let out_dir = Path::new(&out_dir);
    let tonic_services = get_tonic_services(services);
    tonic_prost_build::manual::Builder::new().compile(&tonic_services);

    let content = compile_flatbuffers_tonic_file(&gen_ctx, &types);
    write_tonic_flatbuffers_file(out_dir, &content, &gen_ctx.get_namespace())?;
    Ok(())
}

/// Compile multiple fbs files
pub(crate) fn compile_flatbuffers_tonic_file_list_only<P>(
    path_list: &[P],
) -> Result<(), Box<dyn std::error::Error>>
where
    P: AsRef<Path>,
{
    for path in path_list {
        compile_flatbuffers_tonic_file_only(path.as_ref())?;
    }
    Ok(())
}

fn get_tonic_services(
    services: &[flatbuffers_util::reflect::Service],
) -> Vec<tonic_prost_build::manual::Service> {
    let mut tonic_services = Vec::new();
    for service in services {
        let namespace = service.namespace.as_ref().expect("no namespace found");
        let mut tonic_service = tonic_prost_build::manual::Service::builder()
            .package(namespace)
            .name(&service.name);
        for method in &service.methods {
            let method_name = get_method_name(method);
            let mut tonic_method = tonic_prost_build::manual::Method::builder()
                .name(&method_name)
                .route_name(&method.name)
                .input_type(format!("super::Owned{}", method.request_type()))
                .output_type(format!("super::Owned{}", method.response_type()))
                .codec_path("flatbuffers_tonic::FlatBuffersCodec");
            if method.client_streaming {
                tonic_method = tonic_method.client_streaming();
            }
            if method.server_streaming {
                tonic_method = tonic_method.server_streaming();
            }
            tonic_service = tonic_service.method(tonic_method.build());
        }
        tonic_services.push(tonic_service.build());
    }
    tonic_services
}

fn get_method_name(method: &flatbuffers_util::reflect::Method) -> String {
    // remove the prefix if any, and convert to snake_case
    let name = method.name.rsplit('.').next().unwrap().to_string();
    name.to_snake_case()
}

pub fn compile_types_alias(package: &str, types: &Vec<MessageType>) -> TokenStream {
    let mut content = TokenStream::new();

    for t in types {
        let wrapper_type = quote::format_ident!("Owned{}", t.fb_type);
        let rs_type = quote::format_ident!("{}", t.fb_type);
        let rs_type_mod: syn::Path = syn::parse_str(package).unwrap();
        // add definition
        content.extend(quote! {
            pub struct #wrapper_type(pub flatbuffers_tonic::OwnedFB<#rs_type_mod::#rs_type<'static>>);
            impl #wrapper_type {
                pub fn get_ref<'a>(&'a self) -> #rs_type_mod::#rs_type<'a> {
                    self.0.get_ref()
                }
            }
        });
        // add impl
        content.extend(quote! {
            impl flatbuffers_tonic::OwnedFBCodecable for #wrapper_type {
                fn new_from_bytes(bytes: bytes::Bytes) -> Result<Self, flatbuffers::InvalidFlatbuffer> {
                    let owned = flatbuffers_tonic::OwnedFB::<#rs_type_mod::#rs_type<'static>>::new_from_bytes(bytes)?;
                    Ok(Self(owned))
                }

                fn into_bytes(self) -> bytes::Bytes {
                    self.0.into_bytes()
                }
            }
        });
    }
    content
}

pub fn compile_flatbuffers_tonic_file(
    gen_ctx: &GeneratorContext,
    types: &Vec<MessageType>,
) -> TokenStream {
    let services = gen_ctx.get_services();
    let package_rs = gen_ctx.get_namespace_rs();
    let package = gen_ctx.get_namespace();
    let mut content = TokenStream::new();

    let types_content = compile_types_alias(&package_rs, types);

    content.extend(quote! {
        /// This file is generated by flatbuffers-tonic-build
        /// Do not edit this file manually
    });

    // add flatbuffers includes
    content.extend(quote! {
        // flatbuffers generated code has warnings.
        include!(concat!(env!("OUT_DIR"), "/flatbuffers/", #package, "/mod.rs"));
    });

    // add include for flatbuffers and tonic files
    for service in services {
        let namespace = service.namespace.as_ref().expect("no namespace found");
        let service_name = &service.name;
        content.extend(quote! {
                include!(concat!(env!("OUT_DIR"), "/", #namespace, ".", #service_name, ".rs"));
        });
    }
    // add types content
    content.extend(types_content);
    content
}

fn write_tonic_flatbuffers_file(
    out_dir: &Path,
    content: &TokenStream,
    package_file_suffix: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    use std::fs::File;
    use std::io::Write;

    let dest_path = out_dir.join(format!("flatbuffers_tonic.{package_file_suffix}.rs"));
    let mut f = File::create(&dest_path)?;

    // Parse TokenStream to syn::File and pretty-print
    let syntax_tree: syn::File = syn::parse2(content.clone()).unwrap_or_else(|e| {
        panic!("Failed to parse generated content to syntax tree: {e} : {content}");
    });
    let formatted = prettyplease::unparse(&syntax_tree);
    f.write_all(formatted.as_bytes())?;
    Ok(())
}