tool.rs

use std::collections::HashSet;

use proc_macro2::TokenStream;
use quote::{ToTokens, quote};
use syn::{
    Expr, FnArg, Ident, ItemFn, ItemImpl, MetaList, PatType, Token, Type, Visibility, parse::Parse,
    parse_quote,
};

#[derive(Default)]
struct ToolImplItemAttrs {
    tool_box: Option<Option<Ident>>,
}

impl Parse for ToolImplItemAttrs {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let mut tool_box = None;
        while !input.is_empty() {
            let key: Ident = input.parse()?;
            match key.to_string().as_str() {
                "tool_box" => {
                    tool_box = Some(None);
                    if input.lookahead1().peek(Token![=]) {
                        let value: Ident = input.parse()?;
                        tool_box = Some(Some(value));
                    }
                }
                _ => {
                    return Err(syn::Error::new(key.span(), "unknown attribute"));
                }
            }
            if input.is_empty() {
                break;
            }
            input.parse::<Token![,]>()?;
        }

        Ok(ToolImplItemAttrs { tool_box })
    }
}

#[derive(Default)]
struct ToolFnItemAttrs {
    name: Option<Expr>,
    description: Option<Expr>,
    vis: Option<Visibility>,
}

impl Parse for ToolFnItemAttrs {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let mut name = None;
        let mut description = None;
        let mut vis = None;
        while !input.is_empty() {
            let key: Ident = input.parse()?;
            input.parse::<Token![=]>()?;
            match key.to_string().as_str() {
                "name" => {
                    let value: Expr = input.parse()?;
                    name = Some(value);
                }
                "description" => {
                    let value: Expr = input.parse()?;
                    description = Some(value);
                }
                "vis" => {
                    let value: Visibility = input.parse()?;
                    vis = Some(value);
                }
                _ => {
                    return Err(syn::Error::new(key.span(), "unknown attribute"));
                }
            }
            if input.is_empty() {
                break;
            }
            input.parse::<Token![,]>()?;
        }

        Ok(ToolFnItemAttrs {
            name,
            description,
            vis,
        })
    }
}

struct ToolFnParamAttrs {
    serde_meta: Vec<MetaList>,
    schemars_meta: Vec<MetaList>,
    ident: Ident,
    rust_type: Box<Type>,
}

impl ToTokens for ToolFnParamAttrs {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let ident = &self.ident;
        let rust_type = &self.rust_type;
        let serde_meta = &self.serde_meta;
        let schemars_meta = &self.schemars_meta;
        tokens.extend(quote! {
            #(#[#serde_meta])*
            #(#[#schemars_meta])*
            pub #ident: #rust_type,
        });
    }
}

#[derive(Default)]

enum ToolParams {
    Aggregated {
        rust_type: PatType,
    },
    Params {
        attrs: Vec<ToolFnParamAttrs>,
    },
    #[default]
    NoParam,
}

#[derive(Default)]
struct ToolAttrs {
    fn_item: ToolFnItemAttrs,
    params: ToolParams,
}
const TOOL_IDENT: &str = "tool";
const SERDE_IDENT: &str = "serde";
const SCHEMARS_IDENT: &str = "schemars";
const PARAM_IDENT: &str = "param";
const AGGREGATED_IDENT: &str = "aggr";
const REQ_IDENT: &str = "req";

pub enum ParamMarker {
    Param,
    Aggregated,
}

impl Parse for ParamMarker {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let ident: Ident = input.parse()?;
        match ident.to_string().as_str() {
            PARAM_IDENT => Ok(ParamMarker::Param),
            AGGREGATED_IDENT | REQ_IDENT => Ok(ParamMarker::Aggregated),
            _ => Err(syn::Error::new(ident.span(), "unknown attribute")),
        }
    }
}

pub enum ToolItem {
    Fn(ItemFn),
    Impl(ItemImpl),
}

impl Parse for ToolItem {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let lookahead = input.lookahead1();
        if lookahead.peek(Token![impl]) {
            let item = input.parse::<ItemImpl>()?;
            Ok(ToolItem::Impl(item))
        } else {
            let item = input.parse::<ItemFn>()?;
            Ok(ToolItem::Fn(item))
        }
    }
}

// dispatch impl function item and impl block item
pub(crate) fn tool(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
    let tool_item = syn::parse2::<ToolItem>(input)?;
    match tool_item {
        ToolItem::Fn(item) => tool_fn_item(attr, item),
        ToolItem::Impl(item) => tool_impl_item(attr, item),
    }
}

pub(crate) fn tool_impl_item(attr: TokenStream, mut input: ItemImpl) -> syn::Result<TokenStream> {
    let tool_impl_attr: ToolImplItemAttrs = syn::parse2(attr)?;
    let tool_box_ident = tool_impl_attr.tool_box;
    if input.trait_.is_some() {
        if let Some(ident) = tool_box_ident {
            input.items.push(parse_quote!(
                rmcp::tool_box!(@derive #ident);
            ));
        }
    } else if let Some(ident) = tool_box_ident {
        let mut tool_fn_idents = Vec::new();
        for item in &input.items {
            if let syn::ImplItem::Fn(method) = item {
                for attr in &method.attrs {
                    if attr.path().is_ident(TOOL_IDENT) {
                        tool_fn_idents.push(method.sig.ident.clone());
                    }
                }
            }
        }
        let this_type_ident = &input.self_ty;
        input.items.push(parse_quote!(
            rmcp::tool_box!(#this_type_ident {
                #(#tool_fn_idents),*
            } #ident);
        ));
    }
    Ok(quote! {
        #input
    })
}

pub(crate) fn tool_fn_item(attr: TokenStream, mut input_fn: ItemFn) -> syn::Result<TokenStream> {
    let mut tool_macro_attrs = ToolAttrs::default();
    let args: ToolFnItemAttrs = syn::parse2(attr)?;
    tool_macro_attrs.fn_item = args;
    // let mut fommated_fn_args: Punctuated<FnArg, Comma> = Punctuated::new();
    let mut unextractable_args_indexes = HashSet::new();
    for (index, mut fn_arg) in input_fn.sig.inputs.iter_mut().enumerate() {
        enum Caught {
            Param(ToolFnParamAttrs),
            Aggregated(PatType),
        }
        let mut caught = None;
        match &mut fn_arg {
            FnArg::Receiver(_) => {
                continue;
            }
            FnArg::Typed(pat_type) => {
                let mut serde_metas = Vec::new();
                let mut schemars_metas = Vec::new();
                let mut arg_ident = match pat_type.pat.as_ref() {
                    syn::Pat::Ident(pat_ident) => Some(pat_ident.ident.clone()),
                    _ => None,
                };
                let raw_attrs: Vec<_> = pat_type.attrs.drain(..).collect();
                for attr in raw_attrs {
                    match &attr.meta {
                        syn::Meta::List(meta_list) => {
                            if meta_list.path.is_ident(TOOL_IDENT) {
                                let pat_type = pat_type.clone();
                                let marker = meta_list.parse_args::<ParamMarker>()?;
                                match marker {
                                    ParamMarker::Param => {
                                        let Some(arg_ident) = arg_ident.take() else {
                                            return Err(syn::Error::new(
                                                proc_macro2::Span::call_site(),
                                                "input param must have an ident as name",
                                            ));
                                        };
                                        caught.replace(Caught::Param(ToolFnParamAttrs {
                                            serde_meta: Vec::new(),
                                            schemars_meta: Vec::new(),
                                            ident: arg_ident,
                                            rust_type: pat_type.ty.clone(),
                                        }));
                                    }
                                    ParamMarker::Aggregated => {
                                        caught.replace(Caught::Aggregated(pat_type.clone()));
                                    }
                                }
                            } else if meta_list.path.is_ident(SERDE_IDENT) {
                                serde_metas.push(meta_list.clone());
                            } else if meta_list.path.is_ident(SCHEMARS_IDENT) {
                                schemars_metas.push(meta_list.clone());
                            } else {
                                pat_type.attrs.push(attr);
                            }
                        }
                        _ => {
                            pat_type.attrs.push(attr);
                        }
                    }
                }
                match caught {
                    Some(Caught::Param(mut param)) => {
                        param.serde_meta = serde_metas;
                        param.schemars_meta = schemars_metas;
                        match &mut tool_macro_attrs.params {
                            ToolParams::Params { attrs } => {
                                attrs.push(param);
                            }
                            _ => {
                                tool_macro_attrs.params = ToolParams::Params { attrs: vec![param] };
                            }
                        }
                        unextractable_args_indexes.insert(index);
                    }
                    Some(Caught::Aggregated(rust_type)) => {
                        tool_macro_attrs.params = ToolParams::Aggregated { rust_type };
                        unextractable_args_indexes.insert(index);
                    }
                    None => {}
                }
            }
        }
    }

    // input_fn.sig.inputs = fommated_fn_args;
    let name = if let Some(expr) = tool_macro_attrs.fn_item.name {
        expr
    } else {
        let fn_name = &input_fn.sig.ident;
        parse_quote! {
            stringify!(#fn_name)
        }
    };
    let tool_attr_fn_ident = Ident::new(
        &format!("{}_tool_attr", input_fn.sig.ident),
        proc_macro2::Span::call_site(),
    );

    // generate get tool attr function
    let tool_attr_fn = {
        let description = if let Some(expr) = tool_macro_attrs.fn_item.description {
            expr
        } else {
            parse_quote! {
                ""
            }
        };
        let schema = match &tool_macro_attrs.params {
            ToolParams::Aggregated { rust_type } => {
                let ty = &rust_type.ty;
                let schema = quote! {
                    rmcp::handler::server::tool::cached_schema_for_type::<#ty>()
                };
                schema
            }
            ToolParams::Params { attrs, .. } => {
                let (param_type, temp_param_type_name) =
                    create_request_type(attrs, input_fn.sig.ident.to_string());
                let schema = quote! {
                    {
                        #param_type
                        rmcp::handler::server::tool::cached_schema_for_type::<#temp_param_type_name>()
                    }
                };
                schema
            }
            ToolParams::NoParam => {
                quote! {
                    rmcp::handler::server::tool::cached_schema_for_type::<rmcp::model::EmptyObject>()
                }
            }
        };
        let input_fn_attrs = &input_fn.attrs;
        let input_fn_vis = &input_fn.vis;
        quote! {
            #(#input_fn_attrs)*
            #input_fn_vis fn #tool_attr_fn_ident() -> rmcp::model::Tool {
                rmcp::model::Tool {
                    name: #name.into(),
                    description: Some(#description.into()),
                    input_schema: #schema.into(),
                    annotations: None
                }
            }
        }
    };

    // generate wrapped tool function
    let tool_call_fn = {
        // wrapper function have the same sig:
        // async fn #tool_tool_call(context: rmcp::handler::server::tool::ToolCallContext<'_, Self>)
        //      -> std::result::Result<rmcp::model::CallToolResult, rmcp::Error>
        //
        // and the block part should be like:
        // {
        //      use rmcp::handler::server::tool::*;
        //      let (t0, context) = <T0>::from_tool_call_context_part(context)?;
        //      let (t1, context) = <T1>::from_tool_call_context_part(context)?;
        //      ...
        //      let (tn, context) = <Tn>::from_tool_call_context_part(context)?;
        //      // for params
        //      ... expand helper types here
        //      let (__rmcp_tool_req, context) = rmcp::model::JsonObject::from_tool_call_context_part(context)?;
        //      let __#TOOL_ToolCallParam { param_0, param_1, param_2, .. } = parse_json_object(__rmcp_tool_req)?;
        //      // for aggr
        //      let (Parameters(aggr), context) = <Parameters<AggrType>>::from_tool_call_context_part(context)?;
        //      Self::#tool_ident(to, param_0, t1, param_1, ..., param_2, tn, aggr).await.into_call_tool_result()
        //
        // }
        //
        //
        //

        // for receiver type, name it as __rmcp_tool_receiver
        let is_async = input_fn.sig.asyncness.is_some();
        let receiver_ident = || Ident::new("__rmcp_tool_receiver", proc_macro2::Span::call_site());
        // generate the extraction part for trivial args
        let trivial_args = input_fn
            .sig
            .inputs
            .iter()
            .enumerate()
            .filter_map(|(index, arg)| {
                if unextractable_args_indexes.contains(&index) {
                    None
                } else {
                    // get ident/type pair
                    let line = match arg {
                        FnArg::Typed(pat_type) => {
                            let pat = &pat_type.pat;
                            let ty = &pat_type.ty;
                            quote! {
                                let (#pat, context) = <#ty>::from_tool_call_context_part(context)?;
                            }
                        }
                        FnArg::Receiver(r) => {
                            let ty = r.ty.clone();
                            let pat = receiver_ident();
                            quote! {
                                let  (#pat, context) = <#ty>::from_tool_call_context_part(context)?;
                            }
                        }
                    };
                    Some(line)
                }
            });
        let trivial_arg_extraction_part = quote! {
            #(#trivial_args)*
        };
        let processed_arg_extraction_part = match &mut tool_macro_attrs.params {
            ToolParams::Aggregated { rust_type } => {
                let PatType { pat, ty, .. } = rust_type;
                quote! {
                    let (Parameters(#pat), context) = <Parameters<#ty>>::from_tool_call_context_part(context)?;
                }
            }
            ToolParams::Params { attrs } => {
                let (param_type, temp_param_type_name) =
                    create_request_type(attrs, input_fn.sig.ident.to_string());

                let params_ident = attrs.iter().map(|attr| &attr.ident).collect::<Vec<_>>();
                quote! {
                    #param_type
                    let (__rmcp_tool_req, context) = rmcp::model::JsonObject::from_tool_call_context_part(context)?;
                    let #temp_param_type_name {
                        #(#params_ident,)*
                    } = parse_json_object(__rmcp_tool_req)?;
                }
            }
            ToolParams::NoParam => {
                quote! {}
            }
        };
        // generate the execution part
        // has receiver?
        let params = &input_fn
            .sig
            .inputs
            .iter()
            .map(|fn_arg| match fn_arg {
                FnArg::Receiver(_) => {
                    let pat = receiver_ident();
                    quote! { #pat }
                }
                FnArg::Typed(pat_type) => {
                    let pat = &pat_type.pat.clone();
                    quote! { #pat }
                }
            })
            .collect::<Vec<_>>();
        let raw_fn_ident = &input_fn.sig.ident;
        let call = if is_async {
            quote! {
                Self::#raw_fn_ident(#(#params),*).await.into_call_tool_result()
            }
        } else {
            quote! {
                Self::#raw_fn_ident(#(#params),*).into_call_tool_result()
            }
        };
        // assemble the whole function
        let tool_call_fn_ident = Ident::new(
            &format!("{}_tool_call", input_fn.sig.ident),
            proc_macro2::Span::call_site(),
        );
        let raw_fn_vis = tool_macro_attrs
            .fn_item
            .vis
            .as_ref()
            .unwrap_or(&input_fn.vis);
        let raw_fn_attr = &input_fn
            .attrs
            .iter()
            .filter(|attr| !attr.path().is_ident(TOOL_IDENT))
            .collect::<Vec<_>>();
        quote! {
            #(#raw_fn_attr)*
            #raw_fn_vis async fn #tool_call_fn_ident(context: rmcp::handler::server::tool::ToolCallContext<'_, Self>)
                -> std::result::Result<rmcp::model::CallToolResult, rmcp::Error> {
                use rmcp::handler::server::tool::*;
                #trivial_arg_extraction_part
                #processed_arg_extraction_part
                #call
            }
        }
    };
    Ok(quote! {
        #tool_attr_fn
        #tool_call_fn
        #input_fn
    })
}

fn create_request_type(attrs: &[ToolFnParamAttrs], tool_name: String) -> (TokenStream, Ident) {
    let pascal_case_tool_name = tool_name.to_ascii_uppercase();
    let temp_param_type_name = Ident::new(
        &format!("__{pascal_case_tool_name}ToolCallParam",),
        proc_macro2::Span::call_site(),
    );
    (
        quote! {
            use rmcp::{serde, schemars};
            #[derive(serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
            pub struct #temp_param_type_name {
                #(#attrs)*
            }
        },
        temp_param_type_name,
    )
}

#[cfg(test)]
mod test {
    use super::*;
    #[test]
    fn test_tool_sync_macro() -> syn::Result<()> {
        let attr = quote! {
            name = "test_tool",
            description = "test tool",
            vis =
        };
        let input = quote! {
            fn sum(&self, #[tool(aggr)] req: StructRequest) -> Result<CallToolResult, McpError> {
                Ok(CallToolResult::success(vec![Content::text((req.a + req.b).to_string())]))
            }
        };
        let input = tool(attr, input)?;

        println!("input: {:#}", input);
        Ok(())
    }

    #[test]
    fn test_trait_tool_macro() -> syn::Result<()> {
        let attr = quote! {};
        let input = quote! {
            impl ServerHandler for Calculator {
                tool_box!(@derive);
                fn get_info(&self) -> ServerInfo {
                    ServerInfo {
                        instructions: Some("A simple calculator".into()),
                        ..Default::default()
                    }
                }
            }
        };
        let input = tool(attr, input)?;

        println!("input: {:#}", input);
        Ok(())
    }
}