//! rapace-macros: Proc macros for rapace RPC.
//!
//! Provides `#[rapace::service]` which generates:
//! - Client stubs with async methods
//! - Server dispatch by method_id
//!
//! # Method ID Generation
//!
//! Method IDs are generated by hashing "ServiceName.method_name" using FNV-1a,
//! then truncating to 32 bits. This ensures unique IDs across services without
//! requiring a central registry or manual assignment.

use heck::{ToShoutySnakeCase, ToSnakeCase};
use proc_macro::TokenStream;
use proc_macro_crate::{FoundCrate, crate_name};
use proc_macro2::{Ident, Span, TokenStream as TokenStream2, TokenTree};
use quote::{format_ident, quote};

mod parser;

use parser::{Error as MacroError, ParsedTrait, join_doc_lines, parse_trait};

/// Compute a method ID by hashing "ServiceName.method_name" using FNV-1a.
///
/// This generates globally unique method IDs without requiring sequential
/// assignment or a central registry. The hash is truncated to 32 bits.
fn compute_method_id(service_name: &str, method_name: &str) -> u32 {
    // FNV-1a hash constants
    const FNV_OFFSET: u64 = 0xcbf29ce484222325;
    const FNV_PRIME: u64 = 0x100000001b3;

    let mut hash = FNV_OFFSET;

    // Hash "ServiceName.method_name"
    for byte in service_name.bytes() {
        hash ^= byte as u64;
        hash = hash.wrapping_mul(FNV_PRIME);
    }
    // Hash the dot separator
    hash ^= b'.' as u64;
    hash = hash.wrapping_mul(FNV_PRIME);
    // Hash method name
    for byte in method_name.bytes() {
        hash ^= byte as u64;
        hash = hash.wrapping_mul(FNV_PRIME);
    }

    // Truncate to 32 bits - we XOR the high and low halves to preserve entropy
    ((hash >> 32) ^ hash) as u32
}

/// Generates RPC client and server from a trait definition.
///
/// # Example
///
/// ```ignore
/// #[rapace::service]
/// trait Calculator {
///     async fn add(&self, a: i32, b: i32) -> i32;
/// }
///
/// // Generated:
/// // - CalculatorClient<T: Transport> with async fn add(&self, a: i32, b: i32) -> Result<i32, RpcError>
/// // - CalculatorServer<S: Calculator> with dispatch method
/// ```
///
/// # Streaming RPCs
///
/// For server-streaming, return `Streaming<T>`:
///
/// ```ignore
/// use rapace_core::Streaming;
///
/// #[rapace::service]
/// trait RangeService {
///     async fn range(&self, n: u32) -> Streaming<u32>;
/// }
/// ```
///
/// The client method becomes:
/// `async fn range(&self, n: u32) -> Result<Streaming<u32>, RpcError>`
#[proc_macro_attribute]
pub fn service(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let trait_tokens = TokenStream2::from(item.clone());

    let parsed_trait = match parse_trait(&trait_tokens) {
        Ok(parsed) => parsed,
        Err(err) => return err.to_compile_error().into(),
    };

    match generate_service(&parsed_trait, &trait_tokens) {
        Ok(tokens) => tokens.into(),
        Err(err) => err.to_compile_error().into(),
    }
}

fn generate_service(
    input: &ParsedTrait,
    trait_tokens: &TokenStream2,
) -> Result<TokenStream2, MacroError> {
    let trait_name = &input.ident;
    let trait_name_str = trait_name.to_string();
    let trait_snake = trait_name_str.to_snake_case();
    let trait_shouty = trait_name_str.to_shouty_snake_case();
    let vis = &input.vis_tokens;

    // Detect the rapace crate name
    // Try to find `rapace` first (the facade crate).
    // If not found, check if we're IN the rapace crate itself.
    // If neither, check if rapace_core is available (for internal crates).
    let rapace_crate = match crate_name("rapace") {
        Ok(FoundCrate::Itself) => quote!(rapace),
        Ok(FoundCrate::Name(name)) => {
            let ident = Ident::new(&name, Span::call_site());
            quote!(#ident)
        }
        Err(_) => {
            // rapace not found - check if this is an internal crate with direct dependencies
            if crate_name("rapace_core").is_ok() {
                // We have rapace_core - this is likely an internal crate
                // Create a local rapace module that re-exports what we need
                return Err(MacroError::new(
                    Span::call_site(),
                    "Internal crates using rapace_macros must add `rapace` as a dependency, \
                     or you can create a facade module. See rapace-testkit for an example.",
                ));
            } else {
                return Err(MacroError::new(
                    Span::call_site(),
                    "rapace crate not found in dependencies. Add `rapace = \"...\"` to your Cargo.toml",
                ));
            }
        }
    };

    // Capture trait doc comments
    let trait_doc = join_doc_lines(&input.doc_lines);

    let methods: Vec<MethodInfo> = input
        .methods
        .iter()
        .map(MethodInfo::try_from_parsed)
        .collect::<Result<_, _>>()?;

    let client_name = format_ident!("{}Client", trait_name);
    let server_name = format_ident!("{}Server", trait_name);

    // Generate client methods with hashed method IDs
    let client_methods_hardcoded = methods.iter().map(|m| {
        let method_id = compute_method_id(&trait_name_str, &m.name.to_string());
        generate_client_method(m, method_id, &rapace_crate)
    });

    // Generate client methods that use stored method IDs from registry
    let client_methods_registry = methods
        .iter()
        .enumerate()
        .map(|(idx, m)| generate_client_method_registry(m, idx, &rapace_crate));

    // Generate server dispatch arms (for unary and error fallback)
    let dispatch_arms = methods.iter().map(|m| {
        let method_id = compute_method_id(&trait_name_str, &m.name.to_string());
        generate_dispatch_arm(m, method_id, &rapace_crate)
    });

    // Generate streaming dispatch arms
    let mut streaming_method_matches = Vec::new();
    let streaming_dispatch_arms = methods.iter().map(|m| {
        let method_id = compute_method_id(&trait_name_str, &m.name.to_string());
        generate_streaming_dispatch_arm(m, method_id, &rapace_crate)
    });
    for m in methods.iter() {
        if matches!(m.kind, MethodKind::ServerStreaming { .. }) {
            let method_id = compute_method_id(&trait_name_str, &m.name.to_string());
            streaming_method_matches.push(quote! { #method_id => true, });
        }
    }

    // Generate method ID constants
    let method_id_consts = methods.iter().map(|m| {
        let method_id = compute_method_id(&trait_name_str, &m.name.to_string());
        let method_shouty = m.name.to_string().to_shouty_snake_case();
        let const_name = format_ident!("{}_METHOD_ID_{}", trait_shouty, method_shouty);
        quote! {
            #vis const #const_name: u32 = #method_id;
        }
    });

    // Generate registry registration code
    let register_fn_name = format_ident!("{}_register", trait_snake);
    let register_fn = generate_register_fn(
        &trait_name_str,
        &trait_doc,
        &methods,
        &rapace_crate,
        &register_fn_name,
        vis,
    );

    // Generate registry-aware client struct and constructor
    let registry_client_name = format_ident!("{}RegistryClient", trait_name);
    let method_id_fields: Vec<_> = methods
        .iter()
        .map(|m| {
            let field_name = format_ident!("{}_method_id", m.name);
            quote! { #field_name: u32 }
        })
        .collect();
    let method_id_lookups: Vec<_> = methods
        .iter()
        .map(|m| {
            let field_name = format_ident!("{}_method_id", m.name);
            let method_name = m.name.to_string();
            quote! {
                #field_name: registry.resolve_method_id(#trait_name_str, #method_name)
                    .expect(concat!("method ", #method_name, " not found in registry"))
                    .0
            }
        })
        .collect();

    let expanded = quote! {
        // Keep the original trait
        #trait_tokens

        #(#method_id_consts)*

        #register_fn

        /// Client stub for the #trait_name service.
        ///
        /// This client uses hardcoded method IDs (1, 2, ...) and expects an
        /// [`Arc<RpcSession<T>>`](::std::sync::Arc) whose
        /// [`run`](::#rapace_crate::rapace_core::RpcSession::run) task is already
        /// running. Construct sessions with [`RpcSession::with_channel_start`](::#rapace_crate::rapace_core::RpcSession::with_channel_start) to
        /// coordinate odd/even channel IDs when both peers initiate RPCs.
        /// For multi-service scenarios where method IDs must be globally unique,
        /// use [`#registry_client_name`] instead.
        ///
        /// # Usage
        ///
        /// ```ignore
        /// let session = Arc::new(RpcSession::new(transport));
        /// tokio::spawn(session.clone().run()); // Start the demux loop
        /// let client = FooClient::new(session);
        /// let result = client.some_method(args).await?;
        /// ```
        #vis struct #client_name<T: ::#rapace_crate::rapace_core::Transport> {
            session: ::std::sync::Arc<::#rapace_crate::rapace_core::RpcSession<T>>,
        }

        impl<T: ::#rapace_crate::rapace_core::Transport + Send + Sync + 'static> #client_name<T> {
            /// Create a new client with the given RPC session.
            ///
            /// Uses hardcoded method IDs (1, 2, ...). For registry-based method IDs,
            /// use [`#registry_client_name::new`] instead.
            ///
            /// The provided session must be shared (`Arc::clone`) with the call site
            /// and have its demux loop (`tokio::spawn(session.clone().run())`) running.
            pub fn new(session: ::std::sync::Arc<::#rapace_crate::rapace_core::RpcSession<T>>) -> Self {
                Self { session }
            }

            /// Get a reference to the underlying session.
            pub fn session(&self) -> &::std::sync::Arc<::#rapace_crate::rapace_core::RpcSession<T>> {
                &self.session
            }

            #(#client_methods_hardcoded)*
        }

        /// Registry-aware client stub for the #trait_name service.
        ///
        /// This client looks up method IDs from a [`ServiceRegistry`] at construction time,
        /// ensuring that method IDs are globally unique across all registered services.
        /// It has the same [`RpcSession`](::#rapace_crate::rapace_core::RpcSession) requirements as [`#client_name`].
        #vis struct #registry_client_name<T: ::#rapace_crate::rapace_core::Transport> {
            session: ::std::sync::Arc<::#rapace_crate::rapace_core::RpcSession<T>>,
            #(pub #method_id_fields,)*
        }

        impl<T: ::#rapace_crate::rapace_core::Transport + Send + Sync + 'static> #registry_client_name<T> {
            /// Create a new registry-aware client.
            ///
            /// Looks up method IDs from the registry. The service must be registered
            /// in the registry before calling this constructor.
            ///
            /// The session's demux loop (`session.run()`) must be running for RPC calls to work.
            ///
            /// # Panics
            ///
            /// Panics if the service or any of its methods are not found in the registry.
            pub fn new(session: ::std::sync::Arc<::#rapace_crate::rapace_core::RpcSession<T>>, registry: &::#rapace_crate::registry::ServiceRegistry) -> Self {
                Self {
                    session,
                    #(#method_id_lookups,)*
                }
            }

            /// Get a reference to the underlying session.
            pub fn session(&self) -> &::std::sync::Arc<::#rapace_crate::rapace_core::RpcSession<T>> {
                &self.session
            }

            #(#client_methods_registry)*
        }

        /// Server dispatcher for the #trait_name service.
        ///
        /// Integrate this with an [`RpcSession`](::#rapace_crate::rapace_core::RpcSession)
        /// by calling [`RpcSession::set_dispatcher`](::#rapace_crate::rapace_core::RpcSession::set_dispatcher)
        /// and forwarding `method_id`/`payload` into [`dispatch`] or [`dispatch_streaming`].
        #vis struct #server_name<S> {
            service: S,
        }

        impl<S: #trait_name + Send + Sync + 'static> #server_name<S> {
            /// Create a new server with the given service implementation.
            pub fn new(service: S) -> Self {
                Self { service }
            }

            /// Serve requests from the transport until the connection closes.
            ///
            /// This is the main server loop. It reads frames from the transport,
            /// dispatches them to the appropriate method, and sends responses.
            ///
            /// # Example
            ///
            /// ```ignore
            /// let server = CalculatorServer::new(CalculatorImpl);
            /// server.serve(transport).await?;
            /// ```
            pub async fn serve<T: ::#rapace_crate::rapace_core::Transport + 'static>(
                self,
                transport: ::std::sync::Arc<T>,
            ) -> ::std::result::Result<(), ::#rapace_crate::rapace_core::RpcError> {
                ::#rapace_crate::tracing::debug!("serve: entering loop, waiting for requests");
                loop {
                    // Receive next request frame
                    let request = match transport.recv_frame().await {
                        Ok(frame) => {
                            ::#rapace_crate::tracing::debug!(
                                method_id = frame.desc.method_id,
                                channel_id = frame.desc.channel_id,
                                flags = ?frame.desc.flags,
                                payload_len = frame.payload.len(),
                                "serve: received frame"
                            );
                            frame
                        }
                        Err(::#rapace_crate::rapace_core::TransportError::Closed) => {
                            ::#rapace_crate::tracing::debug!("serve: transport closed");
                            // Connection closed gracefully
                            return Ok(());
                        }
                        Err(e) => {
                            ::#rapace_crate::tracing::error!(?e, "serve: transport error");
                            return Err(::#rapace_crate::rapace_core::RpcError::Transport(e));
                        }
                    };

                    // Skip non-data frames (control frames, etc.)
                    if !request.desc.flags.contains(::#rapace_crate::rapace_core::FrameFlags::DATA) {
                        ::#rapace_crate::tracing::debug!("serve: skipping non-DATA frame");
                        continue;
                    }

                    // Dispatch the request
                    ::#rapace_crate::tracing::debug!(
                        method_id = request.desc.method_id,
                        channel_id = request.desc.channel_id,
                        "serve: dispatching to dispatch_streaming"
                    );
                    if let Err(e) = self.dispatch_streaming(
                        request.desc.method_id,
                        request.desc.channel_id,
                        request.payload,
                        transport.as_ref(),
                    ).await {
                        ::#rapace_crate::tracing::error!(?e, "serve: dispatch_streaming returned error");
                        // Send error response
                        let mut desc = ::#rapace_crate::rapace_core::MsgDescHot::new();
                        desc.channel_id = request.desc.channel_id;
                        desc.flags = ::#rapace_crate::rapace_core::FrameFlags::ERROR | ::#rapace_crate::rapace_core::FrameFlags::EOS;

                        // Encode error: [code: u32 LE][message_len: u32 LE][message bytes]
                        let (code, message): (u32, ::std::string::String) = match &e {
                            ::#rapace_crate::rapace_core::RpcError::Status { code, message } => (*code as u32, message.clone()),
                            ::#rapace_crate::rapace_core::RpcError::Transport(_) => (::#rapace_crate::rapace_core::ErrorCode::Internal as u32, "transport error".into()),
                            ::#rapace_crate::rapace_core::RpcError::Cancelled => (::#rapace_crate::rapace_core::ErrorCode::Cancelled as u32, "cancelled".into()),
                            ::#rapace_crate::rapace_core::RpcError::DeadlineExceeded => (::#rapace_crate::rapace_core::ErrorCode::DeadlineExceeded as u32, "deadline exceeded".into()),
                        };
                        let mut err_bytes = ::std::vec::Vec::with_capacity(8 + message.len());
                        err_bytes.extend_from_slice(&code.to_le_bytes());
                        err_bytes.extend_from_slice(&(message.len() as u32).to_le_bytes());
                        err_bytes.extend_from_slice(message.as_bytes());

                        let frame = ::#rapace_crate::rapace_core::Frame::with_payload(desc, err_bytes);
                        let _ = transport.send_frame(&frame).await;
                    }
                }
            }

            /// Serve a single request from the transport.
            ///
            /// This is useful for testing or when you want to handle each request
            /// individually.
            pub async fn serve_one<T: #rapace_crate::rapace_core::Transport + 'static>(
                &self,
                transport: &T,
            ) -> ::std::result::Result<(), #rapace_crate::rapace_core::RpcError> {
                // Receive next request frame
                let request = transport.recv_frame().await
                    .map_err(#rapace_crate::rapace_core::RpcError::Transport)?;

                // Skip non-data frames
                if !request.desc.flags.contains(#rapace_crate::rapace_core::FrameFlags::DATA) {
                    return Ok(());
                }

                // Dispatch the request
                self.dispatch_streaming(
                    request.desc.method_id,
                    request.desc.channel_id,
                    request.payload,
                    transport,
                ).await
            }

            /// Dispatch a request frame to the appropriate method.
            ///
            /// Returns a response frame on success for unary methods.
            /// For streaming methods, use `dispatch_streaming` instead.
            pub async fn dispatch(
                &self,
                method_id: u32,
                request_payload: &[u8],
            ) -> ::std::result::Result<#rapace_crate::rapace_core::Frame, #rapace_crate::rapace_core::RpcError> {
                match method_id {
                    #(#dispatch_arms)*
                    _ => Err(#rapace_crate::rapace_core::RpcError::Status {
                        code: #rapace_crate::rapace_core::ErrorCode::Unimplemented,
                        message: ::std::format!("unknown method_id: {}", method_id),
                    }),
                }
            }

            /// Dispatch a streaming request to the appropriate method.
            ///
            /// The method sends frames via the provided transport.
            pub async fn dispatch_streaming<T: #rapace_crate::rapace_core::Transport + 'static>(
                &self,
                method_id: u32,
                channel_id: u32,
                request_payload: &[u8],
                transport: &T,
            ) -> ::std::result::Result<(), #rapace_crate::rapace_core::RpcError> {
                #rapace_crate::tracing::debug!(method_id, channel_id, "dispatch_streaming: entered");
                match method_id {
                    #(#streaming_dispatch_arms)*
                    _ => Err(#rapace_crate::rapace_core::RpcError::Status {
                        code: #rapace_crate::rapace_core::ErrorCode::Unimplemented,
                        message: ::std::format!("unknown method_id: {}", method_id),
                    }),
                }
            }
        }
    };

    Ok(expanded)
}

/// Method kind: unary or server-streaming.
#[derive(Clone, Debug)]
#[allow(clippy::large_enum_variant)]
enum MethodKind {
    /// Unary RPC: single request, single response.
    Unary,
    /// Server-streaming: single request, returns `Streaming<T>`.
    ServerStreaming {
        /// The type T in `Streaming<T>`.
        item_type: TokenStream2,
    },
}

struct MethodInfo {
    name: Ident,
    args: Vec<(Ident, TokenStream2)>, // (name, type) pairs, excluding &self
    return_type: TokenStream2,
    kind: MethodKind,
    doc: String,
}

impl MethodInfo {
    fn try_from_parsed(method: &parser::ParsedMethod) -> Result<Self, MacroError> {
        let doc = join_doc_lines(&method.doc_lines);

        let args = method
            .args
            .iter()
            .map(|arg| (arg.name.clone(), arg.ty.clone()))
            .collect();

        let return_type = method.return_type.clone();
        let kind = if let Some(item_type) = extract_streaming_return_type(&return_type) {
            MethodKind::ServerStreaming { item_type }
        } else {
            MethodKind::Unary
        };

        Ok(Self {
            name: method.name.clone(),
            args,
            return_type,
            kind,
            doc,
        })
    }
}

fn generate_client_method(
    method: &MethodInfo,
    method_id: u32,
    rapace_crate: &TokenStream2,
) -> TokenStream2 {
    match &method.kind {
        MethodKind::Unary => generate_client_method_unary(method, method_id, rapace_crate),
        MethodKind::ServerStreaming { item_type } => {
            generate_client_method_server_streaming(method, method_id, item_type, rapace_crate)
        }
    }
}

fn generate_client_method_registry(
    method: &MethodInfo,
    method_index: usize,
    rapace_crate: &TokenStream2,
) -> TokenStream2 {
    match &method.kind {
        MethodKind::Unary => {
            generate_client_method_unary_registry(method, method_index, rapace_crate)
        }
        MethodKind::ServerStreaming { item_type } => {
            generate_client_method_server_streaming_registry(
                method,
                method_index,
                item_type,
                rapace_crate,
            )
        }
    }
}

fn generate_client_method_unary(
    method: &MethodInfo,
    method_id: u32,
    rapace_crate: &TokenStream2,
) -> TokenStream2 {
    let name = &method.name;
    let return_type = &method.return_type;

    let arg_names: Vec<_> = method.args.iter().map(|(name, _)| name).collect();
    let arg_types: Vec<_> = method.args.iter().map(|(_, ty)| ty).collect();

    // Generate the argument list for the function signature
    let fn_args = arg_names.iter().zip(arg_types.iter()).map(|(name, ty)| {
        quote! { #name: #ty }
    });

    // For encoding, serialize args as a tuple using facet_postcard
    let encode_expr = if arg_names.is_empty() {
        quote! { #rapace_crate::facet_postcard::to_vec(&()).unwrap() }
    } else if arg_names.len() == 1 {
        let arg = &arg_names[0];
        quote! { #rapace_crate::facet_postcard::to_vec(&#arg).unwrap() }
    } else {
        quote! { #rapace_crate::facet_postcard::to_vec(&(#(#arg_names.clone()),*)).unwrap() }
    };

    quote! {
        /// Call the #name method on the remote service.
        pub async fn #name(&self, #(#fn_args),*) -> ::std::result::Result<#return_type, #rapace_crate::rapace_core::RpcError> {
            use #rapace_crate::rapace_core::FrameFlags;

            // Encode request using facet_postcard
            let request_bytes: ::std::vec::Vec<u8> = #encode_expr;

            // Call via session
            let channel_id = self.session.next_channel_id();
            let response = self.session.call(channel_id, #method_id, request_bytes).await?;

            // Check for error flag
            if response.flags.contains(FrameFlags::ERROR) {
                return Err(#rapace_crate::rapace_core::parse_error_payload(&response.payload));
            }

            // Decode response using facet_postcard
            let result: #return_type = #rapace_crate::facet_postcard::from_slice(&response.payload)
                .map_err(|e| #rapace_crate::rapace_core::RpcError::Status {
                    code: #rapace_crate::rapace_core::ErrorCode::Internal,
                    message: ::std::format!("decode error: {:?}", e),
                })?;

            Ok(result)
        }
    }
}

fn extract_streaming_return_type(ty: &TokenStream2) -> Option<TokenStream2> {
    let tokens: Vec<TokenTree> = ty.clone().into_iter().collect();
    let mut index = 0;
    while index < tokens.len() {
        match &tokens[index] {
            TokenTree::Ident(ident) if ident == "Streaming" => {
                let mut search = index + 1;
                while search < tokens.len() {
                    match &tokens[search] {
                        TokenTree::Punct(p) if p.as_char() == '<' => {
                            let inner = collect_generic_tokens(&tokens, search)?;
                            return select_stream_item_type(inner);
                        }
                        TokenTree::Punct(p) if p.as_char() == ':' => {
                            search += 1;
                            continue;
                        }
                        _ => break,
                    }
                }
            }
            _ => {}
        }
        index += 1;
    }
    None
}

fn collect_generic_tokens(tokens: &[TokenTree], start: usize) -> Option<TokenStream2> {
    let mut depth = 0usize;
    let mut inner = TokenStream2::new();
    let mut i = start;
    while i < tokens.len() {
        match &tokens[i] {
            TokenTree::Punct(p) if p.as_char() == '<' => {
                depth += 1;
                if depth > 1 {
                    inner.extend(std::iter::once(tokens[i].clone()));
                }
            }
            TokenTree::Punct(p) if p.as_char() == '>' => {
                if depth == 0 {
                    return None;
                }
                depth -= 1;
                if depth == 0 {
                    return Some(inner);
                }
                inner.extend(std::iter::once(tokens[i].clone()));
            }
            other => {
                if depth >= 1 {
                    inner.extend(std::iter::once(other.clone()));
                }
            }
        }
        i += 1;
    }
    None
}

fn select_stream_item_type(inner: TokenStream2) -> Option<TokenStream2> {
    let segments = split_top_level(inner, ',');
    for segment in segments.into_iter().rev() {
        let text = segment.to_string();
        if text.trim().is_empty() {
            continue;
        }
        if text.trim_start().starts_with('\'') {
            continue;
        }
        return Some(segment);
    }
    None
}

fn split_top_level(tokens: TokenStream2, delimiter: char) -> Vec<TokenStream2> {
    let mut parts = Vec::new();
    let mut current = TokenStream2::new();
    let mut angle_depth = 0usize;
    for tt in tokens.into_iter() {
        match &tt {
            TokenTree::Punct(p) if p.as_char() == '<' => {
                angle_depth += 1;
                current.extend(std::iter::once(tt));
            }
            TokenTree::Punct(p) if p.as_char() == '>' => {
                angle_depth = angle_depth.saturating_sub(1);
                current.extend(std::iter::once(tt));
            }
            TokenTree::Punct(p) if p.as_char() == delimiter && angle_depth == 0 => {
                parts.push(current);
                current = TokenStream2::new();
                continue;
            }
            _ => current.extend(std::iter::once(tt)),
        }
    }
    parts.push(current);
    parts
}

fn generate_client_method_server_streaming(
    method: &MethodInfo,
    method_id: u32,
    item_type: &TokenStream2,
    rapace_crate: &TokenStream2,
) -> TokenStream2 {
    let name = &method.name;

    let arg_names: Vec<_> = method.args.iter().map(|(name, _)| name).collect();
    let arg_types: Vec<_> = method.args.iter().map(|(_, ty)| ty).collect();

    // Generate the argument list for the function signature
    let fn_args = arg_names.iter().zip(arg_types.iter()).map(|(name, ty)| {
        quote! { #name: #ty }
    });

    // For encoding, serialize args as a tuple using facet_postcard
    let encode_expr = if arg_names.is_empty() {
        quote! { #rapace_crate::facet_postcard::to_vec(&()).unwrap() }
    } else if arg_names.len() == 1 {
        let arg = &arg_names[0];
        quote! { #rapace_crate::facet_postcard::to_vec(&#arg).unwrap() }
    } else {
        quote! { #rapace_crate::facet_postcard::to_vec(&(#(#arg_names.clone()),*)).unwrap() }
    };

    quote! {
        /// Call the #name server-streaming method on the remote service.
        ///
        /// Returns a stream that yields items as they arrive from the server.
        /// The stream ends when the server sends EOS, or yields an error if
        /// the server sends an ERROR frame.
        pub async fn #name(&self, #(#fn_args),*) -> ::std::result::Result<#rapace_crate::rapace_core::Streaming<#item_type>, #rapace_crate::rapace_core::RpcError> {
            use #rapace_crate::rapace_core::{ErrorCode, RpcError};

            let request_bytes: ::std::vec::Vec<u8> = #encode_expr;

            // Start the streaming call - this registers a tunnel and sends the request
            let mut rx = self.session
                .start_streaming_call(#method_id, request_bytes)
                .await?;

            // Build a Stream<Item = Result<#item_type, RpcError>> with explicit termination on EOS
            let stream = #rapace_crate::rapace_core::try_stream! {
                while let Some(chunk) = rx.recv().await {
                    // Error chunk - parse and return as error
                    if chunk.is_error {
                        let err = #rapace_crate::rapace_core::parse_error_payload(&chunk.payload);
                        Err(err)?;
                    }

                    // Empty EOS chunk - stream is done
                    if chunk.is_eos && chunk.payload.is_empty() {
                        break;
                    }

                    // DATA chunk (possibly with EOS flag for final item) - deserialize
                    let item: #item_type = #rapace_crate::facet_postcard::from_slice(&chunk.payload)
                        .map_err(|e| RpcError::Status {
                            code: ErrorCode::Internal,
                            message: ::std::format!("decode error: {:?}", e),
                        })?;

                    yield item;
                }
            };

            Ok(::std::boxed::Box::pin(stream))
        }
    }
}

fn generate_dispatch_arm(
    method: &MethodInfo,
    method_id: u32,
    rapace_crate: &TokenStream2,
) -> TokenStream2 {
    match &method.kind {
        MethodKind::Unary => generate_dispatch_arm_unary(method, method_id, rapace_crate),
        MethodKind::ServerStreaming { .. } => {
            // Streaming methods are handled by dispatch_streaming, not dispatch
            // For the dispatch() method, return error for streaming methods
            quote! {
                #method_id => {
                    Err(#rapace_crate::rapace_core::RpcError::Status {
                        code: #rapace_crate::rapace_core::ErrorCode::Internal,
                        message: "streaming method called via unary dispatch".into(),
                    })
                }
            }
        }
    }
}

fn generate_streaming_dispatch_arm(
    method: &MethodInfo,
    method_id: u32,
    rapace_crate: &TokenStream2,
) -> TokenStream2 {
    match &method.kind {
        MethodKind::Unary => {
            // For unary methods in streaming dispatch, call the regular dispatch and send the frame
            let name = &method.name;
            let return_type = &method.return_type;
            let arg_names: Vec<_> = method.args.iter().map(|(name, _)| name).collect();
            let arg_types: Vec<_> = method.args.iter().map(|(_, ty)| ty).collect();

            let decode_and_call = if arg_names.is_empty() {
                quote! {
                    let result: #return_type = self.service.#name().await;
                }
            } else if arg_names.len() == 1 {
                let arg = &arg_names[0];
                let ty = &arg_types[0];
                quote! {
                    let #arg: #ty = #rapace_crate::facet_postcard::from_slice(request_payload)
                        .map_err(|e| #rapace_crate::rapace_core::RpcError::Status {
                            code: #rapace_crate::rapace_core::ErrorCode::InvalidArgument,
                            message: ::std::format!("decode error: {:?}", e),
                        })?;
                    let result: #return_type = self.service.#name(#arg).await;
                }
            } else {
                let tuple_type = quote! { (#(#arg_types),*) };
                quote! {
                    let (#(#arg_names),*): #tuple_type = #rapace_crate::facet_postcard::from_slice(request_payload)
                        .map_err(|e| #rapace_crate::rapace_core::RpcError::Status {
                            code: #rapace_crate::rapace_core::ErrorCode::InvalidArgument,
                            message: ::std::format!("decode error: {:?}", e),
                        })?;
                    let result: #return_type = self.service.#name(#(#arg_names),*).await;
                }
            };

            quote! {
                #method_id => {
                    #decode_and_call

                    // Encode and send response frame
                    let response_bytes: ::std::vec::Vec<u8> = #rapace_crate::facet_postcard::to_vec(&result)
                        .map_err(|e| #rapace_crate::rapace_core::RpcError::Status {
                            code: #rapace_crate::rapace_core::ErrorCode::Internal,
                            message: ::std::format!("encode error: {:?}", e),
                        })?;

                    let mut desc = #rapace_crate::rapace_core::MsgDescHot::new();
                    desc.channel_id = channel_id;
                    desc.flags = #rapace_crate::rapace_core::FrameFlags::DATA | #rapace_crate::rapace_core::FrameFlags::EOS;

                    let frame = if response_bytes.len() <= #rapace_crate::rapace_core::INLINE_PAYLOAD_SIZE {
                        #rapace_crate::rapace_core::Frame::with_inline_payload(desc, &response_bytes)
                            .expect("inline payload should fit")
                    } else {
                        #rapace_crate::rapace_core::Frame::with_payload(desc, response_bytes)
                    };

                    transport.send_frame(&frame).await
                        .map_err(#rapace_crate::rapace_core::RpcError::Transport)?;
                    Ok(())
                }
            }
        }
        MethodKind::ServerStreaming { .. } => {
            generate_streaming_dispatch_arm_server_streaming(method, method_id, rapace_crate)
        }
    }
}

fn generate_streaming_dispatch_arm_server_streaming(
    method: &MethodInfo,
    method_id: u32,
    rapace_crate: &TokenStream2,
) -> TokenStream2 {
    let name = &method.name;
    let arg_names: Vec<_> = method.args.iter().map(|(name, _)| name).collect();
    let arg_types: Vec<_> = method.args.iter().map(|(_, ty)| ty).collect();

    let decode_args = if arg_names.is_empty() {
        quote! {}
    } else if arg_names.len() == 1 {
        let arg = &arg_names[0];
        let ty = &arg_types[0];
        quote! {
            let #arg: #ty = #rapace_crate::facet_postcard::from_slice(request_payload)
                .map_err(|e| #rapace_crate::rapace_core::RpcError::Status {
                    code: #rapace_crate::rapace_core::ErrorCode::InvalidArgument,
                    message: ::std::format!("decode error: {:?}", e),
                })?;
        }
    } else {
        let tuple_type = quote! { (#(#arg_types),*) };
        quote! {
            let (#(#arg_names),*): #tuple_type = #rapace_crate::facet_postcard::from_slice(request_payload)
                .map_err(|e| #rapace_crate::rapace_core::RpcError::Status {
                    code: #rapace_crate::rapace_core::ErrorCode::InvalidArgument,
                    message: ::std::format!("decode error: {:?}", e),
                })?;
        }
    };

    let call_args = if arg_names.is_empty() {
        quote! {}
    } else {
        quote! { #(#arg_names),* }
    };

    quote! {
        #method_id => {
            #decode_args

            // Call the service method to get a stream
            let mut stream = self.service.#name(#call_args).await;

            // Iterate over the stream and send frames
            use #rapace_crate::futures::stream::StreamExt;
            #rapace_crate::tracing::debug!(channel_id, "streaming dispatch: starting to iterate stream");

            loop {
                #rapace_crate::tracing::trace!(channel_id, "streaming dispatch: waiting for next item");
                match stream.next().await {
                    Some(Ok(item)) => {
                        #rapace_crate::tracing::debug!(channel_id, "streaming dispatch: got item, encoding");
                        // Encode item
                        let item_bytes: ::std::vec::Vec<u8> = #rapace_crate::facet_postcard::to_vec(&item)
                            .map_err(|e| #rapace_crate::rapace_core::RpcError::Status {
                                code: #rapace_crate::rapace_core::ErrorCode::Internal,
                                message: ::std::format!("encode error: {:?}", e),
                            })?;

                        // Send DATA frame (not EOS yet)
                        let mut desc = #rapace_crate::rapace_core::MsgDescHot::new();
                        desc.channel_id = channel_id;
                        desc.flags = #rapace_crate::rapace_core::FrameFlags::DATA;

                        let frame = if item_bytes.len() <= #rapace_crate::rapace_core::INLINE_PAYLOAD_SIZE {
                            #rapace_crate::rapace_core::Frame::with_inline_payload(desc, &item_bytes)
                                .expect("inline payload should fit")
                        } else {
                            #rapace_crate::rapace_core::Frame::with_payload(desc, item_bytes)
                        };

                        #rapace_crate::tracing::debug!(channel_id, payload_len = frame.payload().len(), "streaming dispatch: sending DATA frame");
                        transport.send_frame(&frame).await
                            .map_err(#rapace_crate::rapace_core::RpcError::Transport)?;
                        #rapace_crate::tracing::debug!(channel_id, "streaming dispatch: DATA frame sent");
                    }
                    Some(Err(err)) => {
                        #rapace_crate::tracing::warn!(channel_id, ?err, "streaming dispatch: got error from stream");
                        // Send ERROR frame and break
                        let mut desc = #rapace_crate::rapace_core::MsgDescHot::new();
                        desc.channel_id = channel_id;
                        desc.flags = #rapace_crate::rapace_core::FrameFlags::ERROR | #rapace_crate::rapace_core::FrameFlags::EOS;

                        // Encode error: [code: u32 LE][message_len: u32 LE][message bytes]
                        let (code, message): (u32, &str) = match &err {
                            #rapace_crate::rapace_core::RpcError::Status { code, message } => (*code as u32, message.as_str()),
                            #rapace_crate::rapace_core::RpcError::Transport(_) => (#rapace_crate::rapace_core::ErrorCode::Internal as u32, "transport error"),
                            #rapace_crate::rapace_core::RpcError::Cancelled => (#rapace_crate::rapace_core::ErrorCode::Cancelled as u32, "cancelled"),
                            #rapace_crate::rapace_core::RpcError::DeadlineExceeded => (#rapace_crate::rapace_core::ErrorCode::DeadlineExceeded as u32, "deadline exceeded"),
                        };
                        let mut err_bytes = Vec::with_capacity(8 + message.len());
                        err_bytes.extend_from_slice(&code.to_le_bytes());
                        err_bytes.extend_from_slice(&(message.len() as u32).to_le_bytes());
                        err_bytes.extend_from_slice(message.as_bytes());

                        let frame = #rapace_crate::rapace_core::Frame::with_payload(desc, err_bytes);
                        transport.send_frame(&frame).await
                            .map_err(#rapace_crate::rapace_core::RpcError::Transport)?;
                        return Ok(());
                    }
                    None => {
                        #rapace_crate::tracing::debug!(channel_id, "streaming dispatch: stream ended, sending EOS");
                        // Stream is complete: send EOS frame
                        let mut desc = #rapace_crate::rapace_core::MsgDescHot::new();
                        desc.channel_id = channel_id;
                        desc.flags = #rapace_crate::rapace_core::FrameFlags::EOS;
                        let frame = #rapace_crate::rapace_core::Frame::new(desc);
                        transport.send_frame(&frame).await
                            .map_err(#rapace_crate::rapace_core::RpcError::Transport)?;
                        #rapace_crate::tracing::debug!(channel_id, "streaming dispatch: EOS sent, returning");
                        return Ok(());
                    }
                }
            }
        }
    }
}

fn generate_dispatch_arm_unary(
    method: &MethodInfo,
    method_id: u32,
    rapace_crate: &TokenStream2,
) -> TokenStream2 {
    let name = &method.name;
    let return_type = &method.return_type;
    let arg_names: Vec<_> = method.args.iter().map(|(name, _)| name).collect();
    let arg_types: Vec<_> = method.args.iter().map(|(_, ty)| ty).collect();

    // Generate decode expression for args
    let decode_and_call = if arg_names.is_empty() {
        quote! {
            // No arguments to decode
            let result: #return_type = self.service.#name().await;
        }
    } else if arg_names.len() == 1 {
        let arg = &arg_names[0];
        let ty = &arg_types[0];
        quote! {
            let #arg: #ty = #rapace_crate::facet_postcard::from_slice(request_payload)
                .map_err(|e| #rapace_crate::rapace_core::RpcError::Status {
                    code: #rapace_crate::rapace_core::ErrorCode::InvalidArgument,
                    message: ::std::format!("decode error: {:?}", e),
                })?;
            let result: #return_type = self.service.#name(#arg).await;
        }
    } else {
        // Multiple args - decode as tuple
        let tuple_type = quote! { (#(#arg_types),*) };
        quote! {
            let (#(#arg_names),*): #tuple_type = #rapace_crate::facet_postcard::from_slice(request_payload)
                .map_err(|e| #rapace_crate::rapace_core::RpcError::Status {
                    code: #rapace_crate::rapace_core::ErrorCode::InvalidArgument,
                    message: ::std::format!("decode error: {:?}", e),
                })?;
            let result: #return_type = self.service.#name(#(#arg_names),*).await;
        }
    };

    quote! {
        #method_id => {
            #decode_and_call

            // Encode response using facet_postcard
            let response_bytes: ::std::vec::Vec<u8> = #rapace_crate::facet_postcard::to_vec(&result)
                .map_err(|e| #rapace_crate::rapace_core::RpcError::Status {
                    code: #rapace_crate::rapace_core::ErrorCode::Internal,
                    message: ::std::format!("encode error: {:?}", e),
                })?;

            // Build response frame
            let mut desc = #rapace_crate::rapace_core::MsgDescHot::new();
            desc.flags = #rapace_crate::rapace_core::FrameFlags::DATA | #rapace_crate::rapace_core::FrameFlags::EOS;

            let frame = if response_bytes.len() <= #rapace_crate::rapace_core::INLINE_PAYLOAD_SIZE {
                #rapace_crate::rapace_core::Frame::with_inline_payload(desc, &response_bytes)
                    .expect("inline payload should fit")
            } else {
                #rapace_crate::rapace_core::Frame::with_payload(desc, response_bytes)
            };

            Ok(frame)
        }
    }
}

/// Generate the `register` function for service registration.
///
/// This generates a function that registers the service and its methods
/// with a `ServiceRegistry`, capturing request/response schemas via facet.
fn generate_register_fn(
    service_name: &str,
    service_doc: &str,
    methods: &[MethodInfo],
    rapace_crate: &TokenStream2,
    register_fn_name: &Ident,
    vis: &TokenStream2,
) -> TokenStream2 {
    let method_registrations: Vec<TokenStream2> = methods
        .iter()
        .map(|m| {
            let method_name = m.name.to_string();
            let method_doc = &m.doc;
            let arg_types: Vec<_> = m.args.iter().map(|(_, ty)| ty).collect();

            // Generate argument info
            let arg_infos: Vec<TokenStream2> = m
                .args
                .iter()
                .map(|(name, ty)| {
                    let name_str = name.to_string();
                    let type_str = quote!(#ty).to_string();
                    quote! {
                        #rapace_crate::registry::ArgInfo {
                            name: #name_str,
                            type_name: #type_str,
                        }
                    }
                })
                .collect();

            // Request shape: tuple of arg types, or () if no args, or single type if one arg
            let request_shape_expr = if arg_types.is_empty() {
                quote! { <() as #rapace_crate::facet_core::Facet>::SHAPE }
            } else if arg_types.len() == 1 {
                let ty = &arg_types[0];
                quote! { <#ty as #rapace_crate::facet_core::Facet>::SHAPE }
            } else {
                quote! { <(#(#arg_types),*) as #rapace_crate::facet_core::Facet>::SHAPE }
            };

            // Response shape: the return type (or inner type for streaming)
            let response_shape_expr = match &m.kind {
                MethodKind::Unary => {
                    let return_type = &m.return_type;
                    quote! { <#return_type as #rapace_crate::facet_core::Facet>::SHAPE }
                }
                MethodKind::ServerStreaming { item_type } => {
                    quote! { <#item_type as #rapace_crate::facet_core::Facet>::SHAPE }
                }
            };

            // Is this a streaming method?
            let is_streaming = matches!(m.kind, MethodKind::ServerStreaming { .. });

            if is_streaming {
                quote! {
                    builder.add_streaming_method(
                        #method_name,
                        #method_doc,
                        vec![#(#arg_infos),*],
                        #request_shape_expr,
                        #response_shape_expr,
                    );
                }
            } else {
                quote! {
                    builder.add_method(
                        #method_name,
                        #method_doc,
                        vec![#(#arg_infos),*],
                        #request_shape_expr,
                        #response_shape_expr,
                    );
                }
            }
        })
        .collect();

    quote! {
        /// Register this service with a registry.
        ///
        /// This function registers the service and all its methods,
        /// capturing request/response schemas and documentation via facet.
        #vis fn #register_fn_name(registry: &mut #rapace_crate::registry::ServiceRegistry) {
            let mut builder = registry.register_service(#service_name, #service_doc);
            #(#method_registrations)*
            builder.finish();
        }
    }
}

/// Generate a unary client method that uses a stored method ID from the registry.
fn generate_client_method_unary_registry(
    method: &MethodInfo,
    _method_index: usize,
    rapace_crate: &TokenStream2,
) -> TokenStream2 {
    let name = &method.name;
    let return_type = &method.return_type;
    let method_id_field = format_ident!("{}_method_id", name);

    let arg_names: Vec<_> = method.args.iter().map(|(name, _)| name).collect();
    let arg_types: Vec<_> = method.args.iter().map(|(_, ty)| ty).collect();

    let fn_args = arg_names.iter().zip(arg_types.iter()).map(|(name, ty)| {
        quote! { #name: #ty }
    });

    let encode_expr = if arg_names.is_empty() {
        quote! { #rapace_crate::facet_postcard::to_vec(&()).unwrap() }
    } else if arg_names.len() == 1 {
        let arg = &arg_names[0];
        quote! { #rapace_crate::facet_postcard::to_vec(&#arg).unwrap() }
    } else {
        quote! { #rapace_crate::facet_postcard::to_vec(&(#(#arg_names.clone()),*)).unwrap() }
    };

    quote! {
        /// Call the #name method on the remote service.
        pub async fn #name(&self, #(#fn_args),*) -> ::std::result::Result<#return_type, #rapace_crate::rapace_core::RpcError> {
            use #rapace_crate::rapace_core::FrameFlags;

            let request_bytes: ::std::vec::Vec<u8> = #encode_expr;

            // Call via session with registry-assigned method ID
            let channel_id = self.session.next_channel_id();
            let response = self.session.call(channel_id, self.#method_id_field, request_bytes).await?;

            if response.flags.contains(FrameFlags::ERROR) {
                return Err(#rapace_crate::rapace_core::parse_error_payload(&response.payload));
            }

            let result: #return_type = #rapace_crate::facet_postcard::from_slice(&response.payload)
                .map_err(|e| #rapace_crate::rapace_core::RpcError::Status {
                    code: #rapace_crate::rapace_core::ErrorCode::Internal,
                    message: ::std::format!("decode error: {:?}", e),
                })?;

            Ok(result)
        }
    }
}

/// Generate a server-streaming client method that uses a stored method ID from the registry.
fn generate_client_method_server_streaming_registry(
    method: &MethodInfo,
    _method_index: usize,
    item_type: &TokenStream2,
    rapace_crate: &TokenStream2,
) -> TokenStream2 {
    let name = &method.name;
    let method_id_field = format_ident!("{}_method_id", name);

    let arg_names: Vec<_> = method.args.iter().map(|(name, _)| name).collect();
    let arg_types: Vec<_> = method.args.iter().map(|(_, ty)| ty).collect();

    let fn_args = arg_names.iter().zip(arg_types.iter()).map(|(name, ty)| {
        quote! { #name: #ty }
    });

    // For encoding, serialize args as a tuple using facet_postcard
    let encode_expr = if arg_names.is_empty() {
        quote! { #rapace_crate::facet_postcard::to_vec(&()).unwrap() }
    } else if arg_names.len() == 1 {
        let arg = &arg_names[0];
        quote! { #rapace_crate::facet_postcard::to_vec(&#arg).unwrap() }
    } else {
        quote! { #rapace_crate::facet_postcard::to_vec(&(#(#arg_names.clone()),*)).unwrap() }
    };

    quote! {
        /// Call the #name server-streaming method on the remote service.
        ///
        /// Returns a stream that yields items as they arrive from the server.
        /// The stream ends when the server sends EOS, or yields an error if
        /// the server sends an ERROR frame.
        pub async fn #name(&self, #(#fn_args),*) -> ::std::result::Result<#rapace_crate::rapace_core::Streaming<#item_type>, #rapace_crate::rapace_core::RpcError> {
            use #rapace_crate::rapace_core::{ErrorCode, RpcError};

            let request_bytes: ::std::vec::Vec<u8> = #encode_expr;

            // Start the streaming call with registry-assigned method ID
            let mut rx = self.session
                .start_streaming_call(self.#method_id_field, request_bytes)
                .await?;

            // Build a Stream<Item = Result<#item_type, RpcError>> with explicit termination on EOS
            let stream = #rapace_crate::rapace_core::try_stream! {
                while let Some(chunk) = rx.recv().await {
                    // Error chunk - parse and return as error
                    if chunk.is_error {
                        let err = #rapace_crate::rapace_core::parse_error_payload(&chunk.payload);
                        Err(err)?;
                    }

                    // Empty EOS chunk - stream is done
                    if chunk.is_eos && chunk.payload.is_empty() {
                        break;
                    }

                    // DATA chunk (possibly with EOS flag for final item) - deserialize
                    let item: #item_type = #rapace_crate::facet_postcard::from_slice(&chunk.payload)
                        .map_err(|e| RpcError::Status {
                            code: ErrorCode::Internal,
                            message: ::std::format!("decode error: {:?}", e),
                        })?;

                    yield item;
                }
            };

            Ok(::std::boxed::Box::pin(stream))
        }
    }
}
