golem-rust-macro 2.0.0-dev.3

// Copyright 2024-2026 Golem Cloud
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::agentic::helpers::{FunctionOutputInfo, is_static_method};
use crate::agentic::{generic_type_in_agent_method_error, generic_type_in_agent_return_type_error};
use quote::{format_ident, quote};
use syn::spanned::Spanned;
use syn::{FnArg, ItemTrait, Type};

pub fn get_remote_client(
    item_trait: &ItemTrait,
    constructor_data_value_param_defs: &[proc_macro2::TokenStream],
    constructor_data_value_param_idents: &[proc_macro2::Ident],
    constructor_agent_config_param_defs: &[proc_macro2::TokenStream],
    constructor_agent_config_param_idents: &[proc_macro2::Ident],
    agent_type_parameter_names: &[String],
) -> proc_macro2::TokenStream {
    let remote_client_type_name = format_ident!("{}Client", item_trait.ident);

    let type_name = item_trait.ident.to_string();

    let remote_agent_methods_info =
        get_remote_agent_methods_info(item_trait, agent_type_parameter_names);

    let method_names = &remote_agent_methods_info.method_names;

    let methods_impl = remote_agent_methods_info.methods_impl;

    let encode_constructor =
        generate_constructor_data_value_params_encoding(constructor_data_value_param_idents);

    let get_method_ident = if method_names.contains_get() {
        format_ident!("get_")
    } else {
        format_ident!("get")
    };

    let agent_config_params_as_rpc_param = {
        let add_rpc_params_entries = constructor_agent_config_param_idents.iter().map(|param_ident|
            quote! { result.append(&mut ::golem_rust::agentic::IntoRpcConfigParam::into_rpc_param(#param_ident, &[])); }
        );

        quote! {
            &{
                let mut result = Vec::new();
                #(#add_rpc_params_entries)*
                result
            }
        }
    };

    let optional_get_with_config_impl = if !constructor_agent_config_param_defs.is_empty() {
        quote! {
            pub fn get_with_config(#(#constructor_data_value_param_defs,)* #(#constructor_agent_config_param_defs,)*) -> #remote_client_type_name {
                let agent_type =
                   golem_rust::golem_agentic::golem::agent::host::get_agent_type(#type_name).expect("Internal Error: Agent type not registered");

                 #encode_constructor

                 let wasm_rpc = golem_rust::golem_agentic::golem::agent::host::WasmRpc::new(#type_name, &data_value, None, #agent_config_params_as_rpc_param);

                 #remote_client_type_name {
                     agent_type_name: #type_name.to_string(),
                     constructor_data: data_value,
                     phantom_id: None,
                     component_id: agent_type.implemented_by,
                     wasm_rpc,
                 }
            }
        }
    } else {
        quote! {}
    };

    let optional_new_phantom_with_config_impl = if !constructor_agent_config_param_defs.is_empty() {
        quote! {
            pub fn new_phantom_with_config(#(#constructor_data_value_param_defs,)* #(#constructor_agent_config_param_defs,)*) -> #remote_client_type_name {
                let agent_type =
                   golem_rust::golem_agentic::golem::agent::host::get_agent_type(#type_name).expect("Internal Error: Agent type not registered");

                #encode_constructor

                let phantom_uuid = golem_rust::Uuid::new_v4();

                let wasm_rpc = golem_rust::golem_agentic::golem::agent::host::WasmRpc::new(#type_name, &data_value, Some(phantom_uuid.into()), #agent_config_params_as_rpc_param);

                #remote_client_type_name {
                    agent_type_name: #type_name.to_string(),
                    constructor_data: data_value,
                    phantom_id: Some(phantom_uuid),
                    component_id: agent_type.implemented_by,
                    wasm_rpc,
                }
            }
        }
    } else {
        quote! {}
    };

    let optional_get_phantom_with_config_impl = if !constructor_agent_config_param_defs.is_empty() {
        quote! {
            pub fn get_phantom_with_config(phantom_id: golem_rust::Uuid, #(#constructor_data_value_param_defs,)* #(#constructor_agent_config_param_defs,)*) -> #remote_client_type_name {
                let agent_type =
                   golem_rust::golem_agentic::golem::agent::host::get_agent_type(#type_name).expect("Internal Error: Agent type not registered");

                #encode_constructor

                let wasm_rpc = golem_rust::golem_agentic::golem::agent::host::WasmRpc::new(#type_name, &data_value, Some(phantom_id.into()), #agent_config_params_as_rpc_param);

                #remote_client_type_name {
                    agent_type_name: #type_name.to_string(),
                    constructor_data: data_value,
                    phantom_id: Some(phantom_id),
                    component_id: agent_type.implemented_by,
                    wasm_rpc,
                }
            }
        }
    } else {
        quote! {}
    };

    quote! {
        pub struct #remote_client_type_name {
            agent_type_name: String,
            constructor_data: golem_rust::golem_agentic::golem::agent::common::DataValue,
            phantom_id: Option<golem_rust::Uuid>,
            component_id: golem_rust::golem_wasm::ComponentId,
            wasm_rpc: golem_rust::golem_agentic::golem::agent::host::WasmRpc,
        }

        impl #remote_client_type_name {
            pub fn #get_method_ident(#(#constructor_data_value_param_defs,)*) -> #remote_client_type_name {
                let agent_type =
                   golem_rust::golem_agentic::golem::agent::host::get_agent_type(#type_name).expect("Internal Error: Agent type not registered");

                 #encode_constructor

                 let wasm_rpc = golem_rust::golem_agentic::golem::agent::host::WasmRpc::new(#type_name, &data_value, None, &[]);

                 #remote_client_type_name {
                     agent_type_name: #type_name.to_string(),
                     constructor_data: data_value,
                     phantom_id: None,
                     component_id: agent_type.implemented_by,
                     wasm_rpc,
                 }
            }

            #optional_get_with_config_impl

            pub fn new_phantom(#(#constructor_data_value_param_defs,)*) -> #remote_client_type_name {
                let agent_type =
                   golem_rust::golem_agentic::golem::agent::host::get_agent_type(#type_name).expect("Internal Error: Agent type not registered");

                #encode_constructor

                let phantom_uuid = golem_rust::Uuid::new_v4();

                let wasm_rpc = golem_rust::golem_agentic::golem::agent::host::WasmRpc::new(#type_name, &data_value, Some(phantom_uuid.into()), &[]);

                #remote_client_type_name {
                    agent_type_name: #type_name.to_string(),
                    constructor_data: data_value,
                    phantom_id: Some(phantom_uuid),
                    component_id: agent_type.implemented_by,
                    wasm_rpc,
                }
            }

            #optional_new_phantom_with_config_impl

            pub fn get_phantom(phantom_id: golem_rust::Uuid, #(#constructor_data_value_param_defs,)*) -> #remote_client_type_name {
                let agent_type =
                   golem_rust::golem_agentic::golem::agent::host::get_agent_type(#type_name).expect("Internal Error: Agent type not registered");

                #encode_constructor

                let wasm_rpc = golem_rust::golem_agentic::golem::agent::host::WasmRpc::new(#type_name, &data_value, Some(phantom_id.into()), &[]);

                #remote_client_type_name {
                    agent_type_name: #type_name.to_string(),
                    constructor_data: data_value,
                    phantom_id: Some(phantom_id),
                    component_id: agent_type.implemented_by,
                    wasm_rpc,
                }
            }

            #optional_get_phantom_with_config_impl

            pub fn phantom_id(&self) -> Option<golem_rust::Uuid> {
                self.phantom_id
            }

            pub fn get_agent_id(&self) -> String {
                golem_rust::golem_agentic::golem::agent::host::make_agent_id(
                    &self.agent_type_name,
                    &self.constructor_data,
                    self.phantom_id.map(|id| id.into()),
                ).expect("Internal Error: Failed to make agent id")
            }

            #methods_impl
        }
    }
}

fn generate_constructor_data_value_params_encoding(
    param_idents: &[proc_macro2::Ident],
) -> proc_macro2::TokenStream {
    match param_idents.len() {
        0 => quote! {
            let data_value = golem_rust::golem_agentic::golem::agent::common::DataValue::Tuple(vec![]);
        },
        1 => {
            let single_ident = &param_idents[0];
            quote! {
                let data_value = golem_rust::agentic::Schema::to_data_value(#single_ident)
                    .expect("Failed to convert constructor parameter");
            }
        }
        _ => quote! {
            let data_value = golem_rust::golem_agentic::golem::agent::common::DataValue::Tuple(vec![
                #(golem_rust::agentic::Schema::to_element_value(#param_idents)
                    .expect("Failed to convert constructor parameter")),*
            ]);
        },
    }
}

fn get_remote_agent_methods_info(
    tr: &ItemTrait,
    type_parameter_names: &[String],
) -> RemoteAgentMethodsInfo {
    let mut agent_method_names = AgentClientMethodNames::new();

    let method_impls = tr
        .items
        .iter()
        .filter_map(|trait_item| {
            let method = extract_method(trait_item)?;

            if is_static_method(&method.sig) {
                return None;
            }

            if let Err(ts) = validate_return_type(&method.sig, type_parameter_names) {
                return Some(ts);
            }

            if let Err(ts) = validate_input_types(&method.sig, type_parameter_names) {
                return Some(ts);
            }

            let input_defs = collect_input_defs_without_principal(&method.sig);
            let input_idents = collect_input_idents_without_principal(&method.sig);

            let method_name = &method.sig.ident;
            let trigger_name = format_ident!("trigger_{}", method_name);
            let schedule_name = format_ident!("schedule_{}", method_name);
            let schedule_cancelable_name = format_ident!("schedule_cancelable_{}", method_name);

            agent_method_names.extend(vec![
                method_name.to_string(),
                trigger_name.to_string(),
                schedule_name.to_string(),
                schedule_cancelable_name.to_string(),
            ]);

            Some(generate_method_code(
                method_name,
                &trigger_name,
                &schedule_name,
                &schedule_cancelable_name,
                &input_defs,
                &input_idents,
                &method.sig,
            ))
        })
        .collect::<Vec<_>>();

    let code = quote! { #(#method_impls)* };
    RemoteAgentMethodsInfo::new(code, agent_method_names)
}

fn extract_method(item: &syn::TraitItem) -> Option<&syn::TraitItemFn> {
    if let syn::TraitItem::Fn(m) = item {
        Some(m)
    } else {
        None
    }
}

fn validate_return_type(
    sig: &syn::Signature,
    type_params: &[String],
) -> Result<(), proc_macro2::TokenStream> {
    if let syn::ReturnType::Type(_, ty) = &sig.output
        && let syn::Type::Path(path) = &**ty
    {
        let ident = &path.path.segments.last().unwrap().ident;

        if ident == "Self" {
            return Ok(()); // skip Self, still valid
        }

        if type_params.contains(&ident.to_string()) {
            return Err(generic_type_in_agent_return_type_error(
                sig.ident.span(),
                &ident.to_string(),
            ));
        }
    }
    Ok(())
}

fn validate_input_types(
    sig: &syn::Signature,
    type_params: &[String],
) -> Result<(), proc_macro2::TokenStream> {
    for fn_arg in &sig.inputs {
        if let FnArg::Typed(pat_type) = fn_arg
            && let Type::Path(type_path) = &*pat_type.ty
        {
            let type_name = type_path.path.segments.last().unwrap().ident.to_string();
            if type_params.contains(&type_name) {
                return Err(generic_type_in_agent_method_error(
                    pat_type.ty.span(),
                    &type_name,
                ));
            }
        }
    }
    Ok(())
}

fn collect_input_defs_without_principal(sig: &syn::Signature) -> Vec<&syn::FnArg> {
    sig.inputs.iter().filter(|arg| match arg {
        FnArg::Receiver(_) => true,
        FnArg::Typed(pat_type) => !matches!(
            &*pat_type.ty,
            Type::Path(type_path) if type_path.path.segments.last().map(|s| s.ident == "Principal").unwrap_or(false)
        ),
    }).collect()
}

fn collect_input_idents_without_principal(sig: &syn::Signature) -> Vec<syn::Ident> {
    sig.inputs
        .iter()
        .filter_map(|arg| {
            if let FnArg::Typed(pat_type) = arg
                && let syn::Pat::Ident(pat_ident) = &*pat_type.pat
            {
                if let Type::Path(type_path) = &*pat_type.ty
                    && type_path
                        .path
                        .segments
                        .last()
                        .is_some_and(|seg| seg.ident == "Principal")
                {
                    return None;
                }
                return Some(pat_ident.ident.clone());
            }
            None
        })
        .collect()
}

fn generate_input_encoding(input_idents: &[syn::Ident]) -> proc_macro2::TokenStream {
    match input_idents.len() {
        0 => quote! {
            let input = golem_rust::golem_agentic::golem::agent::common::DataValue::Tuple(vec![]);
        },
        1 => {
            let ident = &input_idents[0];
            quote! {
                let input = golem_rust::agentic::Schema::to_data_value(#ident)
                    .expect("Failed to encode parameter");
            }
        }
        _ => quote! {
            let input = golem_rust::golem_agentic::golem::agent::common::DataValue::Tuple(vec![
                #(golem_rust::agentic::Schema::to_element_value(#input_idents)
                    .expect("Failed to encode parameter")),*
            ]);
        },
    }
}

fn generate_method_code(
    method_name: &syn::Ident,
    trigger_name: &syn::Ident,
    schedule_name: &syn::Ident,
    schedule_cancelable_name: &syn::Ident,
    input_defs: &[&syn::FnArg],
    input_idents: &[syn::Ident],
    sig: &syn::Signature,
) -> proc_macro2::TokenStream {
    let remote_method_name = method_name.to_string();
    let remote_token = quote! { #remote_method_name };
    let fn_output_info = FunctionOutputInfo::from_signature(sig);
    let return_type = match &sig.output {
        syn::ReturnType::Type(_, ty) => quote! { #ty },
        syn::ReturnType::Default => quote! { () },
    };
    let process_invoke_result = match &sig.output {
        syn::ReturnType::Type(_, ty) if !fn_output_info.is_unit => {
            quote! {
                <#ty as golem_rust::agentic::Schema>::from_data_value(rpc_result_ok)
                    .expect("Failed to deserialize rpc result to return type")
            }
        }
        _ => quote! {},
    };

    let encode_input = generate_input_encoding(input_idents);

    quote! {
        pub async fn #method_name(#(#input_defs),*) -> #return_type {
            #encode_input

            let rpc_result_future = self.wasm_rpc.async_invoke_and_await(
                #remote_token,
                &input
            );

            let rpc_result: Result<golem_rust::golem_agentic::golem::agent::common::DataValue, golem_rust::golem_agentic::golem::agent::host::RpcError> =
                golem_rust::agentic::await_invoke_result(rpc_result_future).await;

            let rpc_result_ok =
                rpc_result.unwrap_or_else(|e| panic!("rpc call to {} failed: {:?}", #remote_token, e));

            #process_invoke_result
        }

        pub fn #trigger_name(#(#input_defs),*) {
            #encode_input

            let rpc_result: Result<(), golem_rust::golem_agentic::golem::agent::host::RpcError> =
                self.wasm_rpc.invoke(#remote_token, &input);

            rpc_result.unwrap_or_else(|e| panic!("rpc call to trigger {} failed: {:?}", #remote_token, e));
        }

        pub fn #schedule_name(#(#input_defs),*, scheduled_time: golem_rust::wasip2::clocks::wall_clock::Datetime) {
            #encode_input

            self.wasm_rpc.schedule_invocation(
                scheduled_time,
                #remote_token,
                &input
            );
        }

        pub fn #schedule_cancelable_name(#(#input_defs),*, scheduled_time: golem_rust::wasip2::clocks::wall_clock::Datetime) -> golem_rust::golem_agentic::golem::agent::host::CancellationToken {
            #encode_input

            self.wasm_rpc.schedule_cancelable_invocation(
                scheduled_time,
                #remote_token,
                &input
            )
        }
    }
}

struct RemoteAgentMethodsInfo {
    methods_impl: proc_macro2::TokenStream,
    method_names: AgentClientMethodNames,
}

impl RemoteAgentMethodsInfo {
    fn new(methods_impl: proc_macro2::TokenStream, method_names: AgentClientMethodNames) -> Self {
        Self {
            methods_impl,
            method_names,
        }
    }
}

#[derive(Debug)]
struct AgentClientMethodNames {
    method_names: Vec<String>,
}

impl AgentClientMethodNames {
    fn new() -> Self {
        Self {
            method_names: vec![],
        }
    }

    fn extend(&mut self, names: Vec<String>) {
        self.method_names.extend(names);
    }

    fn contains(&self, name: &str) -> bool {
        self.method_names.iter().any(|n| n == name)
    }

    fn contains_get(&self) -> bool {
        self.contains("get")
    }
}