miden-base-macros 0.11.0

use std::{env, fs, path::Path};

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

use crate::{
    boilerplate::runtime_boilerplate,
    util::{generated_wit_folder_at, strip_line_comment},
};

const SCRIPT_PACKAGE_VERSION: &str = "1.0.0";

/// 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_dir = env::var("CARGO_MANIFEST_DIR").map_err(|err| {
        syn::Error::new(error_span, format!("failed to read CARGO_MANIFEST_DIR: {err}"))
    })?;
    let manifest_path = Path::new(&manifest_dir).join("Cargo.toml");
    let manifest_content = fs::read_to_string(&manifest_path).map_err(|err| {
        syn::Error::new(
            error_span,
            format!("failed to read manifest '{}': {err}", manifest_path.display()),
        )
    })?;

    let manifest: Value = manifest_content.parse().map_err(|err| {
        syn::Error::new(
            error_span,
            format!("failed to parse manifest '{}': {err}", manifest_path.display()),
        )
    })?;

    let package_table = manifest
        .get("package")
        .and_then(Value::as_table)
        .ok_or_else(|| syn::Error::new(error_span, "manifest missing [package] table"))?;

    let crate_name = package_table
        .get("name")
        .and_then(Value::as_str)
        .ok_or_else(|| syn::Error::new(error_span, "manifest package missing `name`"))?;

    let component_package = package_table
        .get("metadata")
        .and_then(Value::as_table)
        .and_then(|meta| meta.get("component"))
        .and_then(Value::as_table)
        .and_then(|component| component.get("package"))
        .and_then(Value::as_str)
        .ok_or_else(|| {
            syn::Error::new(error_span, "manifest missing package.metadata.component.package")
        })?;

    let mut imports = collect_script_imports(&manifest_dir, package_table, error_span)?;
    imports.sort();

    let package_with_version = if component_package.contains('@') {
        component_package.to_string()
    } else {
        format!("{component_package}@{SCRIPT_PACKAGE_VERSION}")
    };

    let world_name = format!("{}-world", crate_name.replace('_', "-"));

    let mut wit_source = String::new();
    wit_source.push_str(&format!("package {package_with_version};\n\n"));
    wit_source.push_str(&format!("world {world_name} {{\n"));
    for import in imports {
        wit_source.push_str("    import ");
        wit_source.push_str(&import);
        wit_source.push_str(";\n");
    }
    wit_source.push_str(&format!("    export {export_interface};\n"));
    wit_source.push('}');
    wit_source.push('\n');

    Ok(wit_source)
}

fn collect_script_imports(
    manifest_dir: &str,
    package_table: &toml::value::Table,
    error_span: Span,
) -> Result<Vec<String>, syn::Error> {
    let dependencies = package_table
        .get("metadata")
        .and_then(Value::as_table)
        .and_then(|meta| meta.get("miden"))
        .and_then(Value::as_table)
        .and_then(|miden| miden.get("dependencies"))
        .and_then(Value::as_table);

    let mut imports = Vec::new();

    if let Some(dep_table) = dependencies {
        for (dep_name, dep_value) in dep_table.iter() {
            let dep_config = dep_value.as_table().ok_or_else(|| {
                syn::Error::new(
                    error_span,
                    format!(
                        "dependency '{dep_name}' under package.metadata.miden.dependencies must \
                         be a table"
                    ),
                )
            })?;

            let dependency_path =
                dep_config.get("path").and_then(Value::as_str).ok_or_else(|| {
                    syn::Error::new(
                        error_span,
                        format!(
                            "dependency '{dep_name}' under package.metadata.miden.dependencies is \
                             missing a 'path' entry"
                        ),
                    )
                })?;

            let absolute_path = Path::new(manifest_dir).join(dependency_path);
            let canonical = fs::canonicalize(&absolute_path).map_err(|err| {
                syn::Error::new(
                    error_span,
                    format!(
                        "failed to canonicalize dependency '{dep_name}' path '{}': {err}",
                        absolute_path.display()
                    ),
                )
            })?;

            let dependency_info = parse_dependency_wit(&canonical).map_err(|msg| {
                syn::Error::new(
                    error_span,
                    format!("failed to process WIT for dependency '{dep_name}': {msg}"),
                )
            })?;

            let version_suffix = dependency_info
                .version
                .as_ref()
                .map(|version| format!("@{version}"))
                .unwrap_or_default();

            let interface = dependency_info
                .exports
                .first()
                .ok_or_else(|| {
                    syn::Error::new(
                        error_span,
                        format!(
                            "dependency '{dep_name}' did not export any interfaces in its world \
                             definition"
                        ),
                    )
                })?
                .clone();

            imports.push(format!("{}/{interface}{version_suffix}", dependency_info.package));
        }
    }

    Ok(imports)
}

struct DependencyWit {
    package: String,
    version: Option<String>,
    exports: Vec<String>,
}

fn parse_dependency_wit(root: &Path) -> Result<DependencyWit, String> {
    if root.is_file() {
        return parse_wit_file(root)?.ok_or_else(|| {
            format!("WIT file '{}' does not contain a world export", root.display())
        });
    }

    let default_wit_dir = root.join("wit");
    let generated_wit_dir = generated_wit_folder_at(root)?;
    let wit_dirs = [default_wit_dir, generated_wit_dir];
    for wit_dir in &wit_dirs {
        if !wit_dir.exists() {
            continue;
        }
        let mut entries = fs::read_dir(wit_dir)
            .map_err(|err| format!("failed to read '{}': {err}", wit_dir.display()))?
            .collect::<Result<Vec<_>, _>>()
            .map_err(|err| format!("failed to iterate '{}': {err}", wit_dir.display()))?;

        entries.sort_by_key(|entry| entry.file_name());

        for entry in entries {
            let path = entry.path();
            if path.is_dir() {
                if path.file_name().is_some_and(|name| name == "deps") {
                    continue;
                }
                continue;
            }

            if path.extension().and_then(|ext| ext.to_str()) != Some("wit") {
                continue;
            }

            if let Some(info) = parse_wit_file(&path)? {
                return Ok(info);
            }
        }
    }

    Err(format!("no WIT world definition found in directories '{wit_dirs:?}'"))
}

fn parse_wit_file(path: &Path) -> Result<Option<DependencyWit>, String> {
    let contents = fs::read_to_string(path)
        .map_err(|err| format!("failed to read WIT file '{}': {err}", path.display()))?;

    let (package, version) = match extract_package_identifier(&contents) {
        Some(parts) => parts,
        None => return Ok(None),
    };

    let exports = extract_world_exports(&contents);
    if exports.is_empty() {
        return Ok(None);
    }

    Ok(Some(DependencyWit {
        package,
        version,
        exports,
    }))
}

fn extract_package_identifier(contents: &str) -> Option<(String, Option<String>)> {
    for line in contents.lines() {
        let trimmed = strip_line_comment(line).trim_start();
        if let Some(rest) = trimmed.strip_prefix("package ") {
            let token = rest.trim_end_matches(';').trim();
            if let Some((name, version)) = token.split_once('@') {
                return Some((name.trim().to_string(), Some(version.trim().to_string())));
            }
            return Some((token.to_string(), None));
        }
    }
    None
}

fn extract_world_exports(contents: &str) -> Vec<String> {
    let mut exports = Vec::new();

    for line in contents.lines() {
        let trimmed = strip_line_comment(line).trim();
        if let Some(rest) = trimmed.strip_prefix("export ") {
            let rest = rest.trim_end_matches(';').trim();
            let interface = match rest.split_once(':') {
                Some((_, interface)) => interface.trim(),
                None => rest,
            };
            if !interface.is_empty() {
                exports.push(interface.to_string());
            }
        }
    }

    exports
}