async-nodes-macros 0.1.0

use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{FnArg, ItemFn, Pat, parse_macro_input};

/// Formats standard snake_case function names to PascalCase for the generated Struct.
fn to_pascal_case(s: &str) -> String {
    s.split('_')
        .filter(|w| !w.is_empty())
        .map(|w| {
            let mut c = w.chars();
            c.next().map_or_else(String::new, |first| {
                first.to_uppercase().collect::<String>() + c.as_str()
            })
        })
        .collect()
}

fn is_option_type(ty: &syn::Type) -> bool {
    if let syn::Type::Path(type_path) = ty {
        if let Some(segment) = type_path.path.segments.last() {
            return segment.ident == "Option";
        }
    }
    false
}

/// Helper to detect, extract, and optionally parse a key type from `#[context(KeyType)]`.
fn has_and_remove_context_attr(pat_type: &mut syn::PatType) -> Option<Option<syn::Type>> {
    let mut found = None;
    pat_type.attrs.retain(|attr| {
        if attr.path().is_ident("context") {
            // Check if there is an argument like #[context(HeavyResource)]
            if let syn::Meta::List(meta_list) = &attr.meta {
                if let Ok(key_ty) = meta_list.parse_args::<syn::Type>() {
                    found = Some(Some(key_ty));
                } else {
                    found = Some(None);
                }
            } else {
                found = Some(None);
            }
            false // Discard attribute
        } else {
            true // Keep other attributes
        }
    });
    found
}

/// Unpacks the dependency as a cloned, owned type.
/// Rejects reference parameters to prevent lock-guards from surviving across async suspension points.
fn parse_context_arg(
    arg_name: &syn::Ident,
    arg_ty: &syn::Type,
    key_override: Option<syn::Type>,
) -> Result<proc_macro2::TokenStream, syn::Error> {
    if let syn::Type::Reference(type_ref) = arg_ty {
        return Err(syn::Error::new_spanned(
            type_ref,
            "Context parameters must be owned types (e.g., implementing `Clone` or wrapped in an `Arc`). Borrowed references are not permitted.",
        ));
    }

    let key_ty = key_override.as_ref().unwrap_or(arg_ty);
    Ok(quote! {
        let #arg_name = {
            let guard = run.ctx.get::<#key_ty>().ok_or_else(|| {
                ::async_nodes::error::BspError::ValidationError(format!(
                    "Dependency injection failed: TypeMapKey `{}` not found in BspContext",
                    ::std::any::type_name::<#key_ty>()
                ))
            })?;
            (*guard).clone()
        };
    })
}

fn get_inner_type_if_result(ty: &syn::Type) -> &syn::Type {
    if let syn::Type::Path(type_path) = ty {
        if let Some(segment) = type_path.path.segments.last() {
            if segment.ident == "Result" {
                if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
                    if let Some(syn::GenericArgument::Type(inner_ty)) = args.args.first() {
                        return inner_ty;
                    }
                }
            }
        }
    }
    ty
}

/// Procedural attribute transforming an asynchronous function into an explicit `BspNode`.
#[proc_macro_attribute]
pub fn node(attr: TokenStream, item: TokenStream) -> TokenStream {
    let node_name_override = if !attr.is_empty() {
        let lit = parse_macro_input!(attr as syn::LitStr);
        Some(lit.value())
    } else {
        None
    };

    let mut input_fn = parse_macro_input!(item as ItemFn);
    let fn_name = &input_fn.sig.ident;
    let fn_vis = &input_fn.vis;
    let fn_attrs = &input_fn.attrs;

    let struct_name = format_ident!("{}Node", to_pascal_case(&fn_name.to_string()));
    let ext_trait_name = format_ident!("{struct_name}Ports");
    let node_name_str = node_name_override.unwrap_or_else(|| fn_name.to_string());

    let mut input_ports = Vec::new();
    let mut unpack_inputs = Vec::new();
    let mut context_unpacks = Vec::new();
    let mut call_args = Vec::new();
    let mut port_idx: usize = 0; // Explicitly compile-time usize

    let mut trait_signatures = Vec::new();
    let mut handle_methods = Vec::new();

    for arg in input_fn.sig.inputs.iter_mut() {
        if let FnArg::Typed(pat_type) = arg {
            if let Pat::Ident(pat_ident) = &*pat_type.pat {
                let arg_name = pat_ident.ident.clone();
                let arg_ty = (*pat_type.ty).clone();
                call_args.push(quote! { #arg_name });

                if let Some(key_override) = has_and_remove_context_attr(pat_type) {
                    match parse_context_arg(&arg_name, &arg_ty, key_override) {
                        Ok(tokens) => context_unpacks.push(tokens),
                        Err(err) => return err.into_compile_error().into(),
                    }
                } else {
                    let idx = port_idx;
                    port_idx += 1;

                    trait_signatures.push(quote! {
                        #[must_use]
                        fn #arg_name(&self) -> ::async_nodes::graph::InputPin;
                    });

                    handle_methods.push(quote! {
                        fn #arg_name(&self) -> ::async_nodes::graph::InputPin {
                            ::async_nodes::graph::InputPin {
                                node_id: self.id,
                                port: #idx,
                            }
                        }
                    });

                    input_ports.push(quote! { #idx as usize });

                    if is_option_type(&arg_ty) {
                        unpack_inputs.push(quote! {
                            let #arg_name = {
                                if let Some(val) = run.take_input(#idx as usize) {
                                    <#arg_ty as ::async_nodes::value::BspFromValue>::from_value(val, run.node_id, #idx as usize)?
                                } else {
                                    None
                                }
                            };
                        });
                    } else {
                        unpack_inputs.push(quote! {
                            let #arg_name = {
                                let val = run.take_input(#idx as usize)
                                    .ok_or_else(|| ::async_nodes::error::BspError::UnconnectedPort {
                                        node_id: run.node_id,
                                        port_index: #idx as usize,
                                    })?;
                                <#arg_ty as ::async_nodes::value::BspFromValue>::from_value(val, run.node_id, #idx as usize)?
                            };
                        });
                    }
                }
            }
        }
    }

    // Capture the return type path to delegate port count statically
    let return_type = match &input_fn.sig.output {
        syn::ReturnType::Default => quote! { () },
        syn::ReturnType::Type(_, ty) => quote! { #ty },
    };

    match &input_fn.sig.output {
        syn::ReturnType::Default => {}
        syn::ReturnType::Type(_, ty) => {
            let inner_ty = get_inner_type_if_result(ty);
            if let syn::Type::Tuple(type_tuple) = inner_ty {
                for idx in 0..type_tuple.elems.len() {
                    let method_name = format_ident!("output_{}", idx);
                    trait_signatures.push(quote! {
                        #[must_use]
                        fn #method_name(&self) -> ::async_nodes::graph::OutputPin;
                    });
                    handle_methods.push(quote! {
                        fn #method_name(&self) -> ::async_nodes::graph::OutputPin {
                            ::async_nodes::graph::OutputPin {
                                node_id: self.id,
                                port: #idx,
                            }
                        }
                    });
                }
            } else if let syn::Type::Path(path_type) = inner_ty {
                if !path_type.path.is_ident("()") {
                    trait_signatures.push(quote! {
                        #[must_use]
                        fn output(&self) -> ::async_nodes::graph::OutputPin;
                    });
                    handle_methods.push(quote! {
                        fn output(&self) -> ::async_nodes::graph::OutputPin {
                            ::async_nodes::graph::OutputPin {
                                node_id: self.id,
                                port: 0,
                            }
                        }
                    });
                }
            } else {
                trait_signatures.push(quote! {
                    #[must_use]
                    fn output(&self) -> ::async_nodes::graph::OutputPin;
                });
                handle_methods.push(quote! {
                    fn output(&self) -> ::async_nodes::graph::OutputPin {
                        ::async_nodes::graph::OutputPin {
                            node_id: self.id,
                            port: 0,
                        }
                    }
                });
            }
        }
    }

    let expanded = quote! {
        #(#fn_attrs)*
        #[allow(clippy::unused_async, clippy::needless_pass_by_value)]
        #fn_vis #input_fn

        #[derive(Clone, Default)]
        #fn_vis struct #struct_name;

        impl #struct_name {
            pub fn new() -> Self { Self }
        }

        #fn_vis trait #ext_trait_name {
            #(#trait_signatures)*
        }

        impl #ext_trait_name for ::async_nodes::graph::TypedNodeHandle<#struct_name> {
            #(#handle_methods)*
        }

        #[::async_nodes::async_trait::async_trait]
        impl<S: Send> ::async_nodes::node::BspNode<S> for #struct_name {
            fn name(&self) -> &'static str { #node_name_str }
            fn input_ports(&self) -> ::std::vec::Vec<::async_nodes::graph::PortIndex> { ::std::vec![#(#input_ports),*] }
            fn output_ports(&self) -> ::std::vec::Vec<::async_nodes::graph::PortIndex> { <#return_type as ::async_nodes::value::BspOutputPacker>::output_ports() }
            fn is_dynamic_output(&self) -> bool { <#return_type as ::async_nodes::value::BspOutputPacker>::is_dynamic() }
            fn activation(&self) -> ::async_nodes::node::Activation { ::async_nodes::node::Activation::All }

            async fn execute(
                &self,
                run: &mut ::async_nodes::run::RunContext<'_, S>,
            ) -> ::std::result::Result<::async_nodes::node::NodeOutcome, ::async_nodes::error::BspError> {
                // 1. Unpack all input ports safely using static trait resolution
                #(#context_unpacks)*
                #(#unpack_inputs)*

                // 2. Execute the user's native async logic
                let result = #fn_name(#(#call_args),*).await;

                // 3. Serialize outputs or map execution faults
                <#return_type as ::async_nodes::value::BspOutputPacker>::pack(result, run.node_id, #node_name_str)
            }
        }
    };

    TokenStream::from(expanded)
}

/// Procedural attribute transforming a synchronous, blocking function into a `BspNode`.
#[proc_macro_attribute]
pub fn node_blocking(attr: TokenStream, item: TokenStream) -> TokenStream {
    let node_name_override = if !attr.is_empty() {
        let lit = parse_macro_input!(attr as syn::LitStr);
        Some(lit.value())
    } else {
        None
    };

    let mut input_fn = parse_macro_input!(item as ItemFn);
    let fn_name = &input_fn.sig.ident;
    let fn_vis = &input_fn.vis;
    let fn_attrs = &input_fn.attrs;

    let struct_name = format_ident!("{}Node", to_pascal_case(&fn_name.to_string()));
    let ext_trait_name = format_ident!("{struct_name}Ports");
    let node_name_str = node_name_override.unwrap_or_else(|| fn_name.to_string());

    let mut input_ports = Vec::new();
    let mut unpack_inputs = Vec::new();
    let mut context_unpacks = Vec::new();
    let mut call_args = Vec::new();
    let mut port_idx: usize = 0; // Explicitly compile-time usize

    let mut trait_signatures = Vec::new();
    let mut handle_methods = Vec::new();

    for arg in input_fn.sig.inputs.iter_mut() {
        if let FnArg::Typed(pat_type) = arg {
            if let Pat::Ident(pat_ident) = &*pat_type.pat {
                let arg_name = pat_ident.ident.clone();
                let arg_ty = (*pat_type.ty).clone();
                call_args.push(quote! { #arg_name });

                if let Some(key_override) = has_and_remove_context_attr(pat_type) {
                    match parse_context_arg(&arg_name, &arg_ty, key_override) {
                        Ok(tokens) => context_unpacks.push(tokens),
                        Err(err) => return err.into_compile_error().into(),
                    }
                } else {
                    let idx = port_idx;
                    port_idx += 1;

                    trait_signatures.push(quote! {
                        #[must_use]
                        fn #arg_name(&self) -> ::async_nodes::graph::InputPin;
                    });

                    handle_methods.push(quote! {
                        fn #arg_name(&self) -> ::async_nodes::graph::InputPin {
                            ::async_nodes::graph::InputPin {
                                node_id: self.id,
                                port: #idx,
                            }
                        }
                    });

                    input_ports.push(quote! { #idx as usize });

                    if is_option_type(&arg_ty) {
                        unpack_inputs.push(quote! {
                            let #arg_name = {
                                if let Some(val) = run.take_input(#idx as usize) {
                                    <#arg_ty as ::async_nodes::value::BspFromValue>::from_value(val, run.node_id, #idx as usize)?
                                } else {
                                    None
                                }
                            };
                        });
                    } else {
                        unpack_inputs.push(quote! {
                            let #arg_name = {
                                let val = run.take_input(#idx as usize)
                                    .ok_or_else(|| ::async_nodes::error::BspError::UnconnectedPort {
                                        node_id: run.node_id,
                                        port_index: #idx as usize,
                                    })?;
                                <#arg_ty as ::async_nodes::value::BspFromValue>::from_value(val, run.node_id, #idx as usize)?
                            };
                        });
                    }
                }
            }
        }
    }

    let return_type = match &input_fn.sig.output {
        syn::ReturnType::Default => quote! { () },
        syn::ReturnType::Type(_, ty) => quote! { #ty },
    };

    match &input_fn.sig.output {
        syn::ReturnType::Default => {}
        syn::ReturnType::Type(_, ty) => {
            let inner_ty = get_inner_type_if_result(ty);
            if let syn::Type::Tuple(type_tuple) = inner_ty {
                for idx in 0..type_tuple.elems.len() {
                    let method_name = format_ident!("output_{}", idx);
                    trait_signatures.push(quote! {
                        #[must_use]
                        fn #method_name(&self) -> ::async_nodes::graph::OutputPin;
                    });
                    handle_methods.push(quote! {
                        fn #method_name(&self) -> ::async_nodes::graph::OutputPin {
                            ::async_nodes::graph::OutputPin {
                                node_id: self.id,
                                port: #idx,
                            }
                        }
                    });
                }
            } else if let syn::Type::Path(path_type) = inner_ty {
                if !path_type.path.is_ident("()") {
                    trait_signatures.push(quote! {
                        #[must_use]
                        fn output(&self) -> ::async_nodes::graph::OutputPin;
                    });
                    handle_methods.push(quote! {
                        fn output(&self) -> ::async_nodes::graph::OutputPin {
                            ::async_nodes::graph::OutputPin {
                                node_id: self.id,
                                port: 0,
                            }
                        }
                    });
                }
            } else {
                trait_signatures.push(quote! {
                    #[must_use]
                    fn output(&self) -> ::async_nodes::graph::OutputPin;
                });
                handle_methods.push(quote! {
                    fn output(&self) -> ::async_nodes::graph::OutputPin {
                        ::async_nodes::graph::OutputPin {
                            node_id: self.id,
                            port: 0,
                        }
                    }
                });
            }
        }
    }

    let expanded = quote! {
        #(#fn_attrs)*
        #[allow(clippy::needless_pass_by_value)]
        #fn_vis #input_fn

        #[derive(Clone, Default)]
        #fn_vis struct #struct_name;

        impl #struct_name {
            pub fn new() -> Self { Self }
        }

        #fn_vis trait #ext_trait_name {
            #(#trait_signatures)*
        }

        impl #ext_trait_name for ::async_nodes::graph::TypedNodeHandle<#struct_name> {
            #(#handle_methods)*
        }

        #[::async_nodes::async_trait::async_trait]
        impl<S: Send> ::async_nodes::node::BspNode<S> for #struct_name {
            fn name(&self) -> &'static str { #node_name_str }
            fn input_ports(&self) -> ::std::vec::Vec<::async_nodes::graph::PortIndex> { ::std::vec![#(#input_ports),*] }
            fn output_ports(&self) -> ::std::vec::Vec<::async_nodes::graph::PortIndex> { <#return_type as ::async_nodes::value::BspOutputPacker>::output_ports() }
            fn is_dynamic_output(&self) -> bool { <#return_type as ::async_nodes::value::BspOutputPacker>::is_dynamic() }
            fn activation(&self) -> ::async_nodes::node::Activation { ::async_nodes::node::Activation::All }

            async fn execute(
                &self,
                run: &mut ::async_nodes::run::RunContext<'_, S>,
            ) -> ::std::result::Result<::async_nodes::node::NodeOutcome, ::async_nodes::error::BspError> {
                // 1. Unpack all input ports safely using static trait resolution
                #(#context_unpacks)*
                #(#unpack_inputs)*

                // 2. Offload user's synchronous logic to spawn_blocking
                let result = ::async_nodes::tokio::task::spawn_blocking(move || {
                    #fn_name(#(#call_args),*)
                })
                .await
                .map_err(|err| ::async_nodes::error::BspError::NodeExecutionFailed {
                    node_id: run.node_id,
                    node_name: #node_name_str.to_string(),
                    message: ::std::string::ToString::to_string(&err),
                })?;

                // 3. Serialize outputs or map execution faults
                <#return_type as ::async_nodes::value::BspOutputPacker>::pack(result, run.node_id, #node_name_str)
            }
        }
    };

    TokenStream::from(expanded)
}

/// Custom derive macro providing serialization support to arbitrary developer types.
#[proc_macro_derive(BspType)]
pub fn derive_bsp_type(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as syn::DeriveInput);
    let name = &input.ident;
    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();

    let expanded = quote! {
        impl #impl_generics ::async_nodes::value::BspFromValue for #name #ty_generics #where_clause {
            fn from_value(
                value: ::async_nodes::value::BspValue,
                node_id: ::async_nodes::graph::NodeId,
                port: ::async_nodes::graph::PortIndex,
            ) -> ::std::result::Result<Self, ::async_nodes::error::BspError> {
                let found_type = value.type_name();
                let arc_t = value
                    .downcast::<Self>()
                    .map_err(|_| ::async_nodes::error::BspError::TypeMismatch {
                        node_id,
                        port_index: port,
                        expected: ::std::any::type_name::<Self>(),
                        found: found_type,
                    })?;
                Ok((*arc_t).clone())
            }
        }

        impl #impl_generics ::async_nodes::value::BspToValue for #name #ty_generics #where_clause {
            fn to_value(self) -> ::async_nodes::value::BspValue {
                ::async_nodes::value::BspValue::new(self)
            }
        }

        impl #impl_generics ::async_nodes::value::BspOutputPacker for #name #ty_generics #where_clause {
            fn output_ports() -> ::std::vec::Vec<::async_nodes::graph::PortIndex> {
                ::std::vec![0]
            }
            fn pack(
                self,
                _node_id: ::async_nodes::graph::NodeId,
                _node_name: &str,
            ) -> ::std::result::Result<::async_nodes::node::NodeOutcome, ::async_nodes::error::BspError> {
                let mut outputs = ::indexmap::IndexMap::new();
                outputs.insert(0, ::async_nodes::value::BspToValue::to_value(self));
                ::std::result::Result::Ok(::async_nodes::node::NodeOutcome::Success(outputs))
            }
        }
    };

    TokenStream::from(expanded)
}