miden-base-macros 0.12.0

use proc_macro2::{Literal, Span, TokenStream as TokenStream2};
use quote::quote;
use semver::Version;
use syn::{FnArg, ItemFn, Pat, PatIdent, parse_macro_input, spanned::Spanned};

use crate::{
    boilerplate::runtime_boilerplate,
    wit_builder::WitBuilder,
    wit_world::{ManifestPackage, write_world_block},
};

/// Configuration describing the script macro expansion details.
pub(crate) struct ScriptConfig {
    /// Fully-qualified export interface emitted by the generated WIT world.
    pub export_interface: &'static str,
    /// Fully-qualified path to the guest trait implemented by the generated struct.
    pub guest_trait_path: &'static str,
}

/// Configuration for generating a script guest wrapper.
pub(crate) struct GuestWrapperConfig {
    /// Fully-qualified export interface emitted by the generated WIT world.
    pub export_interface: &'static str,
    /// Fully-qualified path to the guest trait implemented by the generated struct.
    pub guest_trait_path: &'static str,
    /// The name of the generated guest struct.
    pub guest_struct_ident: syn::Ident,
    /// Doc string attached to the generated guest struct.
    pub guest_struct_doc: &'static str,
}

/// Generates the shared script wrapper boilerplate (WIT bindings + export glue + guest entrypoint).
pub(crate) fn expand_guest_wrapper(
    error_span: Span,
    config: GuestWrapperConfig,
    user_items: TokenStream2,
    trait_impls: TokenStream2,
    run_body: TokenStream2,
) -> syn::Result<TokenStream2> {
    let inline_wit = build_script_wit(error_span, config.export_interface)?;
    let inline_literal = Literal::string(&inline_wit);

    let export_path: syn::Path = syn::parse_str(config.guest_trait_path).map_err(|err| {
        syn::Error::new(
            error_span,
            format!("failed to parse guest trait path '{}': {err}", config.guest_trait_path),
        )
    })?;

    let runtime_boilerplate = runtime_boilerplate();
    let guest_struct_ident = config.guest_struct_ident;
    let doc = config.guest_struct_doc;

    Ok(quote! {
        #runtime_boilerplate

        #user_items

        ::miden::generate!(inline = #inline_literal);
        self::bindings::export!(#guest_struct_ident);

        #trait_impls

        // Bring ActiveAccount trait into scope so users can call account.get_id(), etc.
        #[allow(unused_imports)]
        use ::miden::active_account::ActiveAccount as _;

        #[doc = #doc]
        pub struct #guest_struct_ident;

        impl #export_path for #guest_struct_ident {
            fn run(arg: ::miden::Word) {
                #run_body
            }
        }
    })
}

/// Expansion logic used by `#[tx_script]`.
pub(crate) fn expand(
    attr: proc_macro::TokenStream,
    item: proc_macro::TokenStream,
    config: ScriptConfig,
) -> proc_macro::TokenStream {
    if !attr.is_empty() {
        return syn::Error::new(Span::call_site(), "this attribute does not accept arguments")
            .into_compile_error()
            .into();
    }

    let input_fn = parse_macro_input!(item as ItemFn);

    let fn_ident = &input_fn.sig.ident;
    if fn_ident != "run" {
        return syn::Error::new(fn_ident.span(), "this attribute must be applied to `fn run`")
            .into_compile_error()
            .into();
    }

    if input_fn.sig.receiver().is_some() {
        return syn::Error::new(input_fn.sig.span(), "this attribute cannot target methods")
            .into_compile_error()
            .into();
    }

    // Parse the optional second parameter (injected wrapper struct).
    // The trait requires `fn run(arg: Word)`, so if the user declares a second parameter,
    // it will be instantiated via `Default::default()` and passed to the user's function.
    let injected_param = match parse_injected_param(&input_fn) {
        Ok(param) => param,
        Err(err) => return err.into_compile_error().into(),
    };

    let struct_ident = quote::format_ident!("Struct");

    // Generate the call to the user's function, with optional injected parameter.
    // Note: `Account` generated from account components in scripts only implements `ActiveAccount`.
    // `NativeAccount` exposes functions that can be called only by account code.
    let (instantiation, call, trait_impl) = match &injected_param {
        Some((ident, ty)) => (
            quote! { let mut #ident = <#ty as ::core::default::Default>::default(); },
            quote! { #fn_ident(arg, &mut #ident) },
            quote! {
                impl ::miden::active_account::ActiveAccount for #ty {}
            },
        ),
        None => (quote! {}, quote! { #fn_ident(arg) }, quote! {}),
    };

    let expanded = match expand_guest_wrapper(
        Span::call_site(),
        GuestWrapperConfig {
            export_interface: config.export_interface,
            guest_trait_path: config.guest_trait_path,
            guest_struct_ident: struct_ident.clone(),
            guest_struct_doc: "Guest entry point generated by the Miden script attribute.",
        },
        quote! { #input_fn },
        trait_impl,
        quote! {
            #instantiation
            #call;
        },
    ) {
        Ok(tokens) => tokens,
        Err(err) => err.into_compile_error(),
    };

    expanded.into()
}

/// Parses the optional injected parameter from the user's `fn run` signature.
///
/// The trait requires `fn run(arg: Word)`. If the user declares a second parameter,
/// it is treated as an "injected" wrapper struct that will be instantiated via
/// `Default::default()` and passed to the user's function.
///
/// Only up to 2 parameters are supported: `(arg: Word)` or `(arg: Word, account: &mut Account)`.
///
/// Returns `Some((ident, type))` if a second parameter exists, `None` otherwise.
fn parse_injected_param(input_fn: &ItemFn) -> syn::Result<Option<(syn::Ident, syn::Type)>> {
    if input_fn.sig.inputs.is_empty() {
        return Err(syn::Error::new(
            input_fn.sig.span(),
            "fn run requires at least one parameter: (arg: Word) or (arg: Word, account: &mut \
             Account)",
        ));
    }

    if input_fn.sig.inputs.len() > 2 {
        return Err(syn::Error::new(
            input_fn.sig.span(),
            "fn run accepts at most 2 parameters: (arg: Word) or (arg: Word, account: &mut \
             Account)",
        ));
    }

    // Validate the first parameter is `arg: Word`
    let first_arg = input_fn.sig.inputs.first().unwrap();
    match first_arg {
        FnArg::Typed(pat_type) => {
            if !matches!(pat_type.pat.as_ref(), Pat::Ident(_)) {
                return Err(syn::Error::new(
                    pat_type.pat.span(),
                    "first parameter must be a simple identifier (e.g., `arg: Word`)",
                ));
            }
            // Check that the type is `Word`
            if !is_type_named(&pat_type.ty, "Word") {
                return Err(syn::Error::new(
                    pat_type.ty.span(),
                    "first parameter must have type `Word` (e.g., `arg: Word`)",
                ));
            }
        }
        FnArg::Receiver(receiver) => {
            return Err(syn::Error::new(receiver.span(), "unexpected receiver argument"));
        }
    }

    let Some(second_arg) = input_fn.sig.inputs.iter().nth(1) else {
        return Ok(None);
    };

    match second_arg {
        FnArg::Typed(pat_type) => {
            let ident = match pat_type.pat.as_ref() {
                Pat::Ident(PatIdent { ident, .. }) => ident.clone(),
                other => {
                    return Err(syn::Error::new(
                        other.span(),
                        "function arguments must be simple identifiers",
                    ));
                }
            };
            let ty = expect_mut_account_type(&pat_type.ty)?;
            Ok(Some((ident, ty)))
        }
        FnArg::Receiver(receiver) => {
            Err(syn::Error::new(receiver.span(), "unexpected receiver argument"))
        }
    }
}

/// Ensures the type is `&mut Account` (allowing paths like `crate::bindings::Account`) and returns
/// the underlying `Account` type.
fn expect_mut_account_type(ty: &syn::Type) -> syn::Result<syn::Type> {
    let syn::Type::Reference(type_ref) = ty else {
        return Err(syn::Error::new(
            ty.span(),
            "second parameter must be typed as `account: &mut Account`",
        ));
    };
    if type_ref.mutability.is_none() {
        return Err(syn::Error::new(
            ty.span(),
            "second parameter must be typed as `account: &mut Account`",
        ));
    }
    if !is_type_named(&type_ref.elem, "Account") {
        return Err(syn::Error::new(
            ty.span(),
            "second parameter must be typed as `account: &mut Account`",
        ));
    }
    Ok((*type_ref.elem).clone())
}

/// Checks if a type's final path segment matches `name` (allowing module-qualified paths like
/// `miden::Word` or `crate::bindings::Account`).
fn is_type_named(ty: &syn::Type, name: &str) -> bool {
    let syn::Type::Path(type_path) = ty else {
        return false;
    };
    if type_path.qself.is_some() {
        return false;
    }
    type_path
        .path
        .segments
        .last()
        .is_some_and(|seg| seg.ident == name && seg.arguments.is_empty())
}

/// Builds an inlined WIT world definition for a script crate.
pub(crate) fn build_script_wit(
    error_span: Span,
    export_interface: &'static str,
) -> Result<String, syn::Error> {
    let manifest = ManifestPackage::load(error_span)?;
    let crate_name = manifest.crate_name(error_span)?;
    let component_package = manifest.component_package(error_span)?;
    let imports = manifest.collect_miden_dependency_imports(error_span)?;
    let world_name = format!("{}-world", crate_name.replace('_', "-"));
    let exports = [export_interface.to_string()];

    let mut wit = WitBuilder::new("#[tx_script]", component_package, &Version::new(1, 0, 0));
    write_world_block(&mut wit, &world_name, &imports, &exports);

    Ok(wit.finish())
}