olai-codegen 0.0.8

//! Request-builder generation for the low-level client.
//!
//! This is part of the **Generation** stage of the Analysis → Planning → Generation → Output
//! pipeline (see [`super`]). For each method of a service it emits a fluent builder struct (e.g.
//! `CreateCatalogBuilder`) plus the `impl` block that lets callers set optional fields and `await`
//! the request. The builders are written to each service's `builders.rs` and complement the
//! low-level client emitted by [`super::client`]: the client constructs a builder, and the builder
//! drives the actual HTTP call.
//!
//! Builder setters carry the proto field documentation as `#[doc = "..."]` attributes, and
//! required-vs-optional partitioning comes from the per-method
//! [`GenerationPlan`](crate::analysis::GenerationPlan) computed during Planning. The emitted token
//! streams are pretty-printed through `super::format_tokens` before being returned as source
//! strings.

use itertools::Itertools;
use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use syn::Path;

/// Generate a `#[doc = "..."]` attribute for a field or method.
///
/// Uses `doc` if present; otherwise generates a default `"Set <field_name>"` message.
fn field_doc_attr(doc: Option<&str>, field_name: &str) -> TokenStream {
    if let Some(doc) = doc {
        let doc_spaced = format!(" {}", doc.trim_start());
        quote! { #[doc = #doc_spaced] }
    } else {
        let set_msg = format!(" Set {field_name}");
        quote! { #[doc = #set_msg] }
    }
}

use super::format_tokens;
use crate::analysis::{BodyField, RequestParam, RequestType};
use crate::codegen::config::Runtime;
use crate::codegen::{MethodHandler, ServiceHandler};
use crate::parsing::types::{BaseType, RenderContext, unified_to_rust};

/// Generate builder code for all request types in a service, targeting `protocol`.
///
/// The builder layer is shared between REST and ConnectRPC — it only constructs the request proto
/// and calls `client.<method>(&request)`, which both low-level clients expose. `protocol` selects
/// which low-level client type the builders bind to (see
/// [`ServiceHandler::low_level_client_type`]) and which methods are in scope: REST builders are
/// emitted only for methods with an HTTP route; ConnectRPC builders are emitted for every method.
pub(crate) fn generate(
    service: &ServiceHandler<'_>,
    protocol: crate::codegen::ClientProtocol,
) -> crate::error::Result<String> {
    let builder_impls: Vec<_> = service
        .methods()
        .filter(|m| protocol == crate::codegen::ClientProtocol::Connect || m.plan.has_http_route)
        .map(|method| generate_request_builder(method, service, protocol))
        .try_collect()?;

    if builder_impls.is_empty() {
        return Ok(String::new());
    }

    let builder_code = generate_builders_module(service, &builder_impls)?;

    Ok(builder_code)
}

/// Generate the complete builders module
fn generate_builders_module(
    service: &ServiceHandler<'_>,
    builders: &[String],
) -> crate::error::Result<String> {
    let builder_tokens: Vec<TokenStream> = builders
        .iter()
        .map(|b| {
            syn::parse_str::<TokenStream>(b).map_err(|source| crate::error::Error::GeneratedParse {
                tokens: b.clone(),
                source,
            })
        })
        .collect::<crate::error::Result<_>>()?;
    let mod_path: Path = service.models_path();
    let result_path: Path =
        syn::parse_str(&service.config.result_type_path).expect("valid result_type_path");

    // The streaming machinery (`into_stream`, `stream_paginated`, the futures
    // `Stream` traits) is emitted only for `List` methods. Pull its imports in
    // only when the service has at least one, so the imports aren't unused.
    let has_list = service
        .methods()
        .any(|m| matches!(m.plan.request_type, RequestType::List));

    // Boxed future/stream aliases the `IntoFuture` / `into_stream` impls return.
    // Natively they're `Send` (`BoxFuture`/`BoxStream`). In dual-transport (WASM)
    // mode the browser transport's futures are `!Send`, so on wasm32 we alias to
    // the `Local*` (non-Send) variants — keeping one set of impls that compile
    // for both targets.
    let box_aliases = if service.config.dual_transport() {
        let stream_alias = if has_list {
            quote! {
                #[cfg(not(target_arch = "wasm32"))]
                type BoxStr<'a, T> = ::futures::stream::BoxStream<'a, T>;
                #[cfg(target_arch = "wasm32")]
                type BoxStr<'a, T> = ::futures::stream::LocalBoxStream<'a, T>;
            }
        } else {
            quote! {}
        };
        quote! {
            #[cfg(not(target_arch = "wasm32"))]
            type BoxFut<'a, T> = ::futures::future::BoxFuture<'a, T>;
            #[cfg(target_arch = "wasm32")]
            type BoxFut<'a, T> = ::futures::future::LocalBoxFuture<'a, T>;
            #stream_alias
        }
    } else {
        let stream_alias = if has_list {
            quote! { type BoxStr<'a, T> = ::futures::stream::BoxStream<'a, T>; }
        } else {
            quote! {}
        };
        quote! {
            type BoxFut<'a, T> = ::futures::future::BoxFuture<'a, T>;
            #stream_alias
        }
    };
    let stream_helper_imports = if has_list {
        quote! {
            use futures::{TryStreamExt, StreamExt};
            use super::super::stream_paginated;
        }
    } else {
        quote! {}
    };

    let tokens = quote! {
        #![allow(unused_mut)]
        // The `use #mod_path::*` / `super::client::*` globs are convenience-wide;
        // not every builders module references every imported name (more so under
        // buffa). Silence unused-import noise rather than per-name pruning.
        #![allow(unused_imports)]
        #box_aliases
        #stream_helper_imports
        use std::future::IntoFuture;
        use #result_path;
        use #mod_path::*;
        use super::client::*;

        #(#builder_tokens)*
    };

    format_tokens(tokens)
}

/// Generate a builder for a specific request type
fn generate_request_builder(
    method: MethodHandler<'_>,
    service: &ServiceHandler<'_>,
    protocol: crate::codegen::ClientProtocol,
) -> crate::error::Result<String> {
    let builder_ident = method.builder_type();
    let request_type_ident = method.input_type().ok_or_else(|| {
        crate::Error::Build(format!(
            "method `{}` has no input type (Empty input); cannot generate builder",
            method.plan.handler_function_name
        ))
    })?;
    let output_type_ident = method.output_type();
    let client_type_ident = service.low_level_client_type(protocol);
    let method_name = format_ident!("{}", method.plan.handler_function_name);

    // Constructor params: all required parameters (path + required body fields)
    let required_params: Vec<&RequestParam> = method.required_parameters().collect();
    // Optional body fields get with_* setter methods (handles oneof variant expansion)
    let (_, optional_body) = method.split_body_fields();
    // Optional query params also get simple with_* setter methods
    let optional_query: Vec<&RequestParam> = method
        .optional_parameters()
        .filter(|p| matches!(p, RequestParam::Query(_)))
        .collect();

    // The `..Default::default()` spread in the constructor is only meaningful
    // when some field is left unset there: either optional fields (set later via
    // `with_*`) exist, or the buffa runtime's hidden `__buffa_unknown_fields`
    // needs filling. When the constructor already sets every field (a prost
    // request with only required params), the spread trips `clippy::needless_update`.
    let has_optional_fields = !optional_body.is_empty() || !optional_query.is_empty();
    let needs_default_spread =
        has_optional_fields || matches!(method.config.runtime, Runtime::Buffa);

    // Generate constructor
    let constructor = generate_constructor(
        &method,
        &request_type_ident,
        &client_type_ident,
        &required_params,
        needs_default_spread,
    );

    // Generate with_* methods: body fields first (handles oneof variants), then query params
    let with_methods_body = generate_with_methods(&method, &optional_body)?;
    let with_methods_query = optional_query
        .iter()
        .map(|param| generate_simple_with_method(&method, param));
    let with_methods: Vec<_> = with_methods_body
        .into_iter()
        .chain(with_methods_query)
        .collect();

    // Generate an operation-specific doc string before `method` is potentially moved.
    let resource_name = method
        .plan
        .handler_function_name
        .split_once('_')
        .map(|x| x.1)
        .unwrap_or(&method.plan.handler_function_name)
        .replace('_', " ");
    let builder_doc = match method.plan.request_type {
        RequestType::List => format!(" Builder for listing {resource_name}"),
        RequestType::Create => format!(" Builder for creating a {resource_name}"),
        RequestType::Get => format!(" Builder for getting a {resource_name}"),
        RequestType::Update => format!(" Builder for updating a {resource_name}"),
        RequestType::Delete => format!(" Builder for deleting a {resource_name}"),
        RequestType::Custom(_) => format!(" Builder for {resource_name}"),
    };

    let into_stream_impl = if matches!(method.plan.request_type, RequestType::List) {
        Some(generate_into_stream_impl(method, &method_name)?)
    } else {
        None
    };

    // Generate IntoFuture implementation
    let into_future_impl = generate_into_future_impl(
        &builder_ident,
        &client_type_ident,
        output_type_ident.as_ref(),
        &method_name,
    );

    let tokens = quote! {
        #[doc = #builder_doc]
        pub struct #builder_ident {
            client: #client_type_ident,
            request: #request_type_ident,
        }

        impl #builder_ident {
            #constructor

            #(#with_methods)*

            #into_stream_impl
        }

        #into_future_impl
    };

    format_tokens(tokens)
}

/// Generate the constructor for the builder
fn generate_constructor(
    method: &MethodHandler<'_>,
    request_type_ident: &proc_macro2::Ident,
    client_type_ident: &proc_macro2::Ident,
    required_params: &[&RequestParam],
    needs_default_spread: bool,
) -> TokenStream {
    let param_list = required_params.iter().map(|param| {
        let field_ident = param.field_ident();
        let param_type = method.field_type(param.field_type(), RenderContext::Constructor);
        quote! { #field_ident: #param_type }
    });

    let field_assignments = required_params.iter().map(|param| {
        let field_ident = param.field_ident();
        let assignment = method.field_assignment(
            param.field_type(),
            &field_ident,
            &RenderContext::Constructor,
        );
        // When the assignment is just the field ident (no `.into()`/cast), use field shorthand
        // (`field`) instead of `field: field`, which trips `redundant_field_names`.
        if field_ident == assignment.to_string() {
            quote! { #field_ident }
        } else {
            quote! { #field_ident: #assignment }
        }
    });

    let obtain_doc = format!(" Obtain via the corresponding method on `{client_type_ident}`.");
    let default_spread = if needs_default_spread {
        quote! { ..Default::default() }
    } else {
        quote! {}
    };
    quote! {
        #[doc = " Create a new builder instance."]
        #[doc = #obtain_doc]
        pub(crate) fn new(client: #client_type_ident, #(#param_list),*) -> Self {
            let request = #request_type_ident {
                #(#field_assignments,)*
                #default_spread
            };
            Self { client, request }
        }
    }
}

/// Generate with_* methods for optional fields
fn generate_with_methods(
    method: &MethodHandler<'_>,
    optional_fields: &[&BodyField],
) -> crate::error::Result<Vec<TokenStream>> {
    let mut methods = Vec::new();

    // First, generate individual methods for oneof variants
    for field in optional_fields {
        if matches!(field.field_type.base_type, BaseType::OneOf(_))
            && field.oneof_variants.is_some()
        {
            methods.extend(generate_oneof_variant_methods(method, field)?);
        }
    }

    // Then generate regular methods for non-oneof fields
    let regular_methods: Vec<TokenStream> = optional_fields
        .iter()
        .filter(|field| !matches!(field.field_type.base_type, BaseType::OneOf(_)))
        .map(|field| builder_with_impl(method, field))
        .collect();

    methods.extend(regular_methods);
    Ok(methods)
}

fn builder_with_impl(method: &MethodHandler<'_>, field: &BodyField) -> TokenStream {
    let field_ident = format_ident!("{}", field.name);
    let method_name = format_ident!("with_{}", field.name);
    let field_name = &field.name;

    let doc_attr = field_doc_attr(field.documentation.as_deref(), field_name);

    if matches!(field.field_type.base_type, BaseType::Map(_, _)) {
        quote! {
            #doc_attr
            pub fn #method_name<I, K, V>(mut self, #field_ident: I) -> Self
            where
                I: IntoIterator<Item = (K, V)>,
                K: Into<String>,
                V: Into<String>,
            {
                self.request.#field_ident = #field_ident
                    .into_iter()
                    .map(|(k, v)| (k.into(), v.into()))
                    .collect();
                self
            }
        }
    } else if matches!(field.field_type.base_type, BaseType::Enum(_)) {
        let enum_ident = method.field_type(&field.field_type, RenderContext::BuilderMethod);
        let assignment = method.field_assignment(
            &field.field_type,
            &field_ident,
            &RenderContext::BuilderMethod,
        );
        quote! {
            #doc_attr
            pub fn #method_name(mut self, #field_ident: impl Into<Option<#enum_ident>>) -> Self {
                self.request.#field_ident = #assignment;
                self
            }
        }
    } else {
        let field_type = method.field_type(&field.field_type, RenderContext::BuilderMethod);
        if field.repeated {
            let assignment = quote! { #field_ident.into_iter().collect() };
            quote! {
                #doc_attr
                pub fn #method_name<I>(mut self, #field_ident: I) -> Self
                where
                    I: IntoIterator<Item = #field_type>,
                {
                    self.request.#field_ident = #assignment;
                    self
                }
            }
        } else {
            // Catch-all for non-enum, non-repeated fields: singular messages/oneofs
            // (stored as buffa `MessageField<T>` / prost `Option<Box<T>>`) and optional
            // scalars (`Option<T>`). The arg is `impl Into<Option<T>>`.
            //
            // Only buffa's `MessageField` needs the explicit `From<Option<T>>` step
            // (with the intermediate annotated to disambiguate the chained `.into()`);
            // optional scalars and the prost message representation are reached by a
            // plain `Option::into`.
            let is_singular_message = matches!(
                field.field_type.base_type,
                BaseType::Message(_) | BaseType::OneOf(_)
            );
            let assignment =
                if is_singular_message && matches!(method.config.runtime, Runtime::Buffa) {
                    quote! {
                        {
                            let #field_ident: ::core::option::Option<_> = #field_ident.into();
                            buffa::MessageField::from(#field_ident)
                        }
                    }
                } else {
                    quote! { #field_ident.into() }
                };
            quote! {
                #doc_attr
                pub fn #method_name(mut self, #field_ident: impl Into<Option<#field_type>>) -> Self {
                    self.request.#field_ident = #assignment;
                    self
                }
            }
        }
    }
}

/// Generate individual methods for each variant of a oneof field
fn generate_oneof_variant_methods(
    method: &MethodHandler<'_>,
    field: &BodyField,
) -> crate::error::Result<Vec<TokenStream>> {
    // No variants → no methods to emit. Keeps the invariant self-contained rather than relying
    // on the caller having checked `is_some()`.
    let Some(variants) = field.oneof_variants.as_ref() else {
        return Ok(Vec::new());
    };
    let oneof_field_ident = format_ident!("{}", field.name);

    let enum_type_tokens = method.field_type(&field.field_type, RenderContext::BuilderMethod);

    variants
        .iter()
        .map(|variant| {
            let method_name = format_ident!("with_{}", variant.field_name);
            let param_ident = format_ident!("{}", variant.field_name.replace("_", ""));
            let variant_name = format_ident!("{}", variant.variant_name);

            // Derive the Rust parameter type from the UnifiedType abstraction. Hard-fail with
            // context when the derived type string is unparsable rather than silently
            // substituting `String`, which would generate a wrong-typed builder method.
            let rust_type_str = unified_to_rust(&variant.field_type, RenderContext::Parameter);
            let param_type: syn::Type = syn::parse_str(&rust_type_str).map_err(|source| {
                crate::error::Error::InvalidRustPath {
                    path: rust_type_str.clone(),
                    message: format!(
                        "oneof variant `{}` of field `{}`",
                        variant.field_name, field.name
                    ),
                    source,
                }
            })?;

            let doc_attr = field_doc_attr(variant.documentation.as_deref(), &variant.field_name);

            Ok(quote! {
                #doc_attr
                pub fn #method_name(mut self, #param_ident: #param_type) -> Self {
                    self.request.#oneof_field_ident = Some(#enum_type_tokens::#variant_name(#param_ident));
                    self
                }
            })
        })
        .collect()
}

/// Generate IntoFuture implementation
fn generate_into_future_impl(
    builder_ident: &proc_macro2::Ident,
    _client_type_ident: &proc_macro2::Ident,
    output_type_ident: Option<&proc_macro2::Ident>,
    method_name: &proc_macro2::Ident,
) -> TokenStream {
    if let Some(out_ident) = output_type_ident {
        quote! {
            impl IntoFuture for #builder_ident {
                type Output = Result<#out_ident>;
                type IntoFuture = BoxFut<'static, Self::Output>;

                fn into_future(self) -> Self::IntoFuture {
                    let client = self.client;
                    let request = self.request;
                    Box::pin(async move { client.#method_name(&request).await })
                }
            }
        }
    } else {
        quote! {
            impl IntoFuture for #builder_ident {
                type Output = Result<()>;
                type IntoFuture = BoxFut<'static, Self::Output>;

                fn into_future(self) -> Self::IntoFuture {
                    let client = self.client;
                    let request = self.request;
                    Box::pin(async move { client.#method_name(&request).await })
                }
            }
        }
    }
}

/// Generate a simple `with_*` setter for a single optional query parameter.
fn generate_simple_with_method(method: &MethodHandler<'_>, param: &RequestParam) -> TokenStream {
    let field_ident = param.field_ident();
    let method_name = format_ident!("with_{}", param.name());
    let field_name = param.name();

    let doc_attr = field_doc_attr(param.documentation(), field_name);

    let field_type = method.field_type(param.field_type(), RenderContext::BuilderMethod);

    // Enum fields must be cast to i32 at the FFI boundary
    if matches!(param.field_type().base_type, BaseType::Enum(_)) {
        let assignment = method.field_assignment(
            param.field_type(),
            &field_ident,
            &RenderContext::BuilderMethod,
        );
        quote! {
            #doc_attr
            pub fn #method_name(mut self, #field_ident: impl Into<Option<#field_type>>) -> Self {
                self.request.#field_ident = #assignment;
                self
            }
        }
    } else {
        quote! {
            #doc_attr
            pub fn #method_name(mut self, #field_ident: impl Into<Option<#field_type>>) -> Self {
                self.request.#field_ident = #field_ident.into();
                self
            }
        }
    }
}

fn generate_into_stream_impl(
    method: MethodHandler<'_>,
    client_method_name: &proc_macro2::Ident,
) -> crate::error::Result<TokenStream> {
    let items_field = method.list_output_field().ok_or_else(|| {
        crate::Error::Build(format!(
            "List method `{}` has no non-page_token output field; cannot generate into_stream",
            method.plan.handler_function_name
        ))
    })?;
    let item_field_ident = format_ident!("{}", items_field.name);
    let output_type_ident = items_field.unified_type.type_ident();

    // The browser stream is `!Send`, so on wasm32 box it with `boxed_local()`
    // (→ `LocalBoxStream`, the `BoxStr` alias resolves to). Natively, `boxed()`.
    let boxer = if method.config.dual_transport() {
        quote! {
            #[cfg(not(target_arch = "wasm32"))]
            let stream = stream.boxed();
            #[cfg(target_arch = "wasm32")]
            let stream = stream.boxed_local();
            stream
        }
    } else {
        quote! { stream.boxed() }
    };

    Ok(quote! {
        /// Convert paginated request into stream of results
        pub fn into_stream(self) -> BoxStr<'static, Result<#output_type_ident>> {
            let remaining = self.request.max_results;
            let stream = stream_paginated((self, remaining), move |(mut builder, mut remaining), page_token| async move {
                builder.request.page_token = page_token;
                let res = builder.client.#client_method_name(&builder.request).await?;
                if let Some(ref mut rem) = remaining {
                    *rem -= res.#item_field_ident.len() as i32;
                }
                let next_page_token = if remaining.is_some_and(|r| r <= 0) {
                    None
                } else {
                    res.next_page_token.clone()
                };
                Ok((res, (builder, remaining), next_page_token))
            })
            .map_ok(|resp| futures::stream::iter(resp.#item_field_ident.into_iter().map(Ok)))
            .try_flatten();
            #boxer
        }
    })
}

#[cfg(test)]
mod tests {
    use crate::parsing::MessageField;
    use crate::parsing::types::{BaseType, UnifiedType};

    // Note: Tests for enum conversion and field analysis are now integrated into the MethodHandler
    // and would be better tested as part of handler integration tests

    #[test]
    fn test_field_analysis_simulation() {
        // Create a temporary method handler to use the centralized method
        // Note: In practice, this test would be refactored to test the handler methods directly
        let fields = vec![
            MessageField {
                name: "name".to_string(),
                unified_type: UnifiedType {
                    base_type: BaseType::String,
                    is_optional: false,
                    is_repeated: false,
                },
                documentation: None,
                oneof_variants: None,
                field_behavior: vec![],
                is_sensitive: false,
                resource_reference: None,
            },
            MessageField {
                name: "comment".to_string(),
                unified_type: UnifiedType {
                    base_type: BaseType::String,
                    is_optional: true,
                    is_repeated: false,
                },
                documentation: None,
                oneof_variants: None,
                field_behavior: vec![],
                is_sensitive: false,
                resource_reference: None,
            },
            MessageField {
                name: "properties".to_string(),
                unified_type: UnifiedType {
                    base_type: BaseType::Map(
                        Box::new(UnifiedType {
                            base_type: BaseType::String,
                            is_optional: false,
                            is_repeated: false,
                        }),
                        Box::new(UnifiedType {
                            base_type: BaseType::String,
                            is_optional: false,
                            is_repeated: false,
                        }),
                    ),
                    is_optional: true,
                    is_repeated: false,
                },
                documentation: None,
                oneof_variants: None,
                field_behavior: vec![],
                is_sensitive: false,
                resource_reference: None,
            },
        ];

        let mut required = Vec::new();
        let mut optional = Vec::new();

        for field in &fields {
            if field.unified_type.is_optional
                || field.unified_type.is_repeated
                || matches!(
                    field.unified_type.base_type,
                    BaseType::Map(_, _) | BaseType::Message(_) | BaseType::OneOf(_)
                )
            {
                optional.push(field);
            } else {
                required.push(field);
            }
        }

        let (required, optional) = (required, optional);

        assert_eq!(required.len(), 1);
        assert_eq!(required[0].name, "name");

        assert_eq!(optional.len(), 2);
        assert_eq!(optional[0].name, "comment");
        assert_eq!(optional[1].name, "properties");
    }
}