// rhai_macros_derive/src/lib.rs

// We will add our derive macro implementations here.

extern crate proc_macro;

use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::{quote, format_ident, quote_spanned}; 
use syn::{parse_macro_input, Ident, Type, ItemFn, spanned::Spanned, PathArguments, GenericArgument, DeriveInput, Data, LitStr};

// Old ToRhaiMap and FromRhaiMap definitions will be removed from here.
// The export_fn macro definition starts after this.

// Trait definitions removed from here as proc-macro crates cannot export them.
// They should be defined in a regular library crate (e.g., rhai_wrapper or a new rhai_traits crate).

#[proc_macro_attribute]
pub fn export_fn(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let func = parse_macro_input!(item as ItemFn);

    let fn_name = &func.sig.ident;
    let wrapper_fn_name = Ident::new(&format!("rhai_wrapper_{}", fn_name), fn_name.span());
    let original_fn_generics = &func.sig.generics;
    let fn_inputs = &func.sig.inputs;
    // num_expected_args has been removed as Rhai handles arity with typed arguments

    let mut arg_conversions: Vec<TokenStream2> = Vec::new();   // For let converted_arg_0 = ...
    let mut original_fn_call_args: Vec<TokenStream2> = Vec::new(); // For fn_name(converted_arg_0, ...)

    let mut wrapper_args_dynamic_defs: Vec<TokenStream2> = Vec::new(); // Stores definitions like `wrapper_arg_0: ::rhai::Dynamic`

    for (i, fn_arg) in fn_inputs.iter().enumerate() { 
        if let syn::FnArg::Typed(pat_type) = fn_arg { 
            let arg_ident_original_pat = match *pat_type.pat.clone() { 
                syn::Pat::Ident(pat_ident) => pat_ident.ident,
                _ => panic!("#[export_fn] only supports simple identifiers as argument patterns, not complex patterns."),
            };
            let arg_ident_original_pat_str = arg_ident_original_pat.to_string();
            let wrapper_arg_name_ident = Ident::new(&format!("wrapper_arg_{}", i), pat_type.pat.span());
            let wrapper_arg_name_str = wrapper_arg_name_ident.to_string();

            // Populate definitions for the wrapper function signature
            wrapper_args_dynamic_defs.push(quote! { #wrapper_arg_name_ident: ::rhai::Dynamic });

            // Generate a unique ident for the converted argument variable
            let converted_rhai_arg_ident = Ident::new(&format!("converted_rhai_arg_{}", i), pat_type.pat.span());

            let arg_rust_type = &*pat_type.ty; 
            let arg_type_str = quote!(#arg_rust_type).to_string().replace(' ', "");

            if let Type::Reference(ref type_ref) = *pat_type.ty { 
                let is_mutable = type_ref.mutability.is_some();
                let referent_type = &*type_ref.elem; 
                let lock_guard_ident = Ident::new(&format!("lock_guard_{}", i), pat_type.pat.span());

                let conversion_logic = if is_mutable {
                    quote! {
                        let mut #lock_guard_ident = #wrapper_arg_name_ident.write_lock::<#referent_type>()
                            .ok_or_else(|| Box::new(::rhai::EvalAltResult::ErrorMismatchDataType(
                                format!("Argument '{}' (dynamic arg '{}') at index {} must be a mutable reference to {}. Ensure it's not already locked.",
                                    #arg_ident_original_pat_str, #wrapper_arg_name_str, #i, stringify!(#referent_type)
                                ),
                                format!("but found type {}", #wrapper_arg_name_ident.type_name()),
                                ::rhai::Position::NONE
                            )))?;
                        let #converted_rhai_arg_ident = &mut *#lock_guard_ident;
                    }
                } else {
                    quote! {
                        let #lock_guard_ident = #wrapper_arg_name_ident.read_lock::<#referent_type>()
                            .ok_or_else(|| Box::new(::rhai::EvalAltResult::ErrorMismatchDataType(
                                format!("Argument '{}' (dynamic arg '{}') at index {} must be an immutable reference to {}.",
                                    #arg_ident_original_pat_str, #wrapper_arg_name_str, #i, stringify!(#referent_type)
                                ),
                                format!("but found type {}", #wrapper_arg_name_ident.type_name()),
                                ::rhai::Position::NONE
                            )))?;
                        let #converted_rhai_arg_ident = &*#lock_guard_ident;
                    }
                };
                arg_conversions.push(conversion_logic);
                original_fn_call_args.push(quote! { #converted_rhai_arg_ident });
            } else if arg_type_str == "&str" { 
                let str_lock_guard_ident = Ident::new(&format!("str_lock_guard_{}", i), pat_type.pat.span());
                let conversion_logic = quote! {
                    let #str_lock_guard_ident = #wrapper_arg_name_ident.read_lock::<rhai::ImmutableString>()
                        .ok_or_else(|| Box::new(::rhai::EvalAltResult::ErrorMismatchDataType(
                            format!("Argument '{}' (dynamic arg '{}') at index {} must be of type &str (requires read_lock on ImmutableString)",
                                #arg_ident_original_pat_str, #wrapper_arg_name_str, #i
                            ),
                            format!("but found type {}", #wrapper_arg_name_ident.type_name()),
                            ::rhai::Position::NONE
                        )))?;
                     let #converted_rhai_arg_ident = #str_lock_guard_ident.as_str();
                };
                arg_conversions.push(conversion_logic);
                original_fn_call_args.push(quote! { #converted_rhai_arg_ident });
            } else if arg_type_str == "INT" || arg_type_str == "i64" {
                let conversion_logic = quote! {
                    let int_option: Option<::rhai::INT> = #wrapper_arg_name_ident.as_int();
                    let #converted_rhai_arg_ident = int_option
                        .ok_or_else(|| Box::new(::rhai::EvalAltResult::ErrorMismatchDataType(
                            format!("Argument '{}' (dynamic arg '{}') at index {} must be of type i64 (INT)",
                                #arg_ident_original_pat_str, #wrapper_arg_name_str, #i
                            ),
                            format!("but found type {}", #wrapper_arg_name_ident.type_name()),
                            ::rhai::Position::NONE
                        )))?;
                };
                arg_conversions.push(conversion_logic);
                original_fn_call_args.push(quote! { #converted_rhai_arg_ident });
            } else if arg_type_str == "FLOAT" || arg_type_str == "f64" {
                let conversion_logic = quote! {
                    let float_option: Option<::rhai::FLOAT> = #wrapper_arg_name_ident.as_float();
                    let #converted_rhai_arg_ident = float_option
                        .ok_or_else(|| Box::new(::rhai::EvalAltResult::ErrorMismatchDataType(
                            format!("Argument '{}' (dynamic arg '{}') at index {} must be of type f64 (FLOAT)",
                                #arg_ident_original_pat_str, #wrapper_arg_name_str, #i
                            ),
                            format!("but found type {}", #wrapper_arg_name_ident.type_name()),
                            ::rhai::Position::NONE
                        )))?;
                };
                arg_conversions.push(conversion_logic);
                original_fn_call_args.push(quote! { #converted_rhai_arg_ident });
            } else if arg_type_str == "bool" {
                let conversion_logic = quote! {
                    let bool_option: Option<bool> = #wrapper_arg_name_ident.as_bool();
                    let #converted_rhai_arg_ident = bool_option
                        .ok_or_else(|| Box::new(::rhai::EvalAltResult::ErrorMismatchDataType(
                            format!("Argument '{}' (dynamic arg '{}') at index {} must be of type bool",
                                #arg_ident_original_pat_str, #wrapper_arg_name_str, #i
                            ),
                            format!("but found type {}", #wrapper_arg_name_ident.type_name()),
                            ::rhai::Position::NONE
                        )))?;
                };
                arg_conversions.push(conversion_logic);
                original_fn_call_args.push(quote! { #converted_rhai_arg_ident });
            } else if arg_type_str == "String" {
                let conversion_logic = quote! {
                    let #converted_rhai_arg_ident = #wrapper_arg_name_ident.clone().into_string()
                        .map_err(|_| Box::new(::rhai::EvalAltResult::ErrorMismatchDataType(
                            format!("Argument '{}' (dynamic arg '{}') at index {} must be of type String",
                                #arg_ident_original_pat_str, #wrapper_arg_name_str, #i
                            ),
                            format!("but found type {}", #wrapper_arg_name_ident.type_name()),
                            ::rhai::Position::NONE
                        )))?;
                };
                arg_conversions.push(conversion_logic);
                original_fn_call_args.push(quote! { #converted_rhai_arg_ident });
            } else {
                let owned_type_path = &*pat_type.ty; // No deref needed here, quote! can handle &Type
                let arg_type_for_conversion_code = quote!{#owned_type_path};
                let conversion_logic = quote! {
                    let dyn_val_for_cast = #wrapper_arg_name_ident.clone(); // Clone for try_cast
                    let actual_type_name = #wrapper_arg_name_ident.type_name();
                    let #converted_rhai_arg_ident:#arg_type_for_conversion_code = match <#arg_type_for_conversion_code as std::convert::TryFrom<::rhai::Dynamic>>::try_from(dyn_val_for_cast) {
                        Ok(val) => val,
                        Err(_e) => { 
                            let cast_option: Option<#arg_type_for_conversion_code> = #wrapper_arg_name_ident.clone().try_cast::<#arg_type_for_conversion_code>();
                            cast_option
                                .ok_or_else(|| Box::new(::rhai::EvalAltResult::ErrorMismatchDataType(
                                    format!("Argument '{}' (dynamic arg '{}') at index {} must be convertible to type {}. TryFrom and try_cast failed.",
                                        #arg_ident_original_pat_str, #wrapper_arg_name_str, #i, stringify!(#owned_type_path)
                                    ),
                                    format!("but found type {}", actual_type_name),
                                    ::rhai::Position::NONE
                                )))?
                        }
                    };
                };
                arg_conversions.push(conversion_logic);
                original_fn_call_args.push(quote! { #converted_rhai_arg_ident });
            }
        } else {
            // Handle receivers like self, &self, &mut self if necessary
            panic!("#[export_fn] does not support methods with 'self' receivers. Apply to static or free functions.");
        }
    }

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

    // Determine the return type for the wrapper function's Result
    // If original is (), wrapper returns Dynamic (representing unit). Otherwise, original return type.
    let return_type_for_wrapper = if quote!(#return_type).to_string().replace(" ", "") == "()" {
        quote!(::rhai::Dynamic)
    } else {
        quote!(#return_type)
    };

    // How to convert the result of the original function call into the wrapper's Ok Result.
    let return_conversion = if let syn::ReturnType::Type(_, _ty) = &func.sig.output {
        // If original function returns non-(), result is already #return_type_for_wrapper (which is #return_type)
        // So, it can be directly used.
        quote! { Ok(result) }
    } else {
        // If original function returns (), result is (), convert to Dynamic::UNIT for Rhai
        quote! { Ok(().into()) }
    };

    let (impl_generics, _ty_generics, where_clause) = original_fn_generics.split_for_impl();

    let expanded = quote! {
        #func // Re-emit the original function

        // _cx is NativeCallContext. The arguments are now individual ::rhai::Dynamic types.
        pub fn #wrapper_fn_name #impl_generics (_cx: ::rhai::NativeCallContext, #(#wrapper_args_dynamic_defs),*) -> Result<#return_type_for_wrapper, Box<::rhai::EvalAltResult>> #where_clause {
            // Arity check is implicitly handled by Rhai's function registration for typed arguments.
            // If the script calls with wrong number of args, Rhai won't find a matching function signature.
            
            #(#arg_conversions)*
            
            // The following allow might be needed if the original function's result is directly usable
            // without explicit conversion to Dynamic, but the wrapper expects Dynamic for unit returns.
            #[allow(clippy::useless_conversion)] 
            let result = #fn_name(#(#original_fn_call_args),*);
            #return_conversion
        }
    };

    // eprintln!("Generated code for {}:\n{}", stringify!(#fn_name), expanded.to_string());
    proc_macro::TokenStream::from(expanded)
}


#[proc_macro_derive(FromRhaiMap, attributes(rhai_map_field))]
pub fn derive_from_rhai_map(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let name = &input.ident;

    // Helper function to check if a type is Option<T> and extract T
    // Returns (is_option, inner_type_quote_option)
    fn get_option_inner_type(ty: &Type) -> (bool, Option<&Type>) {
        if let Type::Path(type_path) = ty {
            if type_path.path.segments.len() == 1 && type_path.path.segments.first().unwrap().ident == "Option" {
                if let PathArguments::AngleBracketed(params) = &type_path.path.segments.first().unwrap().arguments {
                    if params.args.len() == 1 {
                        if let GenericArgument::Type(inner_ty) = params.args.first().unwrap() {
                            return (true, Some(inner_ty));
                        }
                    }
                }
            }
        }
        (false, None)
    }

    // Helper function to check if a type is Vec<T> and extract T
    // Returns (is_vec, inner_type_quote_option)
    fn get_vec_inner_type(ty: &Type) -> (bool, Option<&Type>) {
        if let Type::Path(type_path) = ty {
            if type_path.path.segments.len() == 1 && type_path.path.segments.first().unwrap().ident == "Vec" {
                if let PathArguments::AngleBracketed(params) = &type_path.path.segments.first().unwrap().arguments {
                    if params.args.len() == 1 {
                        if let GenericArgument::Type(inner_ty) = params.args.first().unwrap() {
                            return (true, Some(inner_ty));
                        }
                    }
                }
            }
        }
        (false, None)
    }

    // Helper to get the string representation of a type, simplified (e.g., "MyStruct", "String", "i64")
    fn get_simple_type_str(ty: &Type) -> String {
        if let Type::Path(type_path) = ty {
            if let Some(segment) = type_path.path.segments.last() {
                return segment.ident.to_string();
            }
        }
        quote!(#ty).to_string().replace(' ', "").replace("::", "_") // fallback, might need refinement
    }

    // Helper to check if a type is a primitive type based on its string representation
    fn is_primitive_type_str(simple_type_str: &str) -> bool {
        ["String", "INT", "i64", "FLOAT", "f64", "bool"].contains(&simple_type_str)
    }

    let fields_data = match &input.data {
        Data::Struct(syn::DataStruct { fields: syn::Fields::Named(fields), .. }) => &fields.named,
        _ => panic!("FromRhaiMapDerive only supports structs with named fields"),
    };

    let mut field_value_declarations = Vec::new();
    let mut struct_field_assignments = Vec::new();

    for field in fields_data.iter() {
        let field_name_ident = field.ident.as_ref().unwrap(); // Changed variable name for clarity
        let field_name_str = field_name_ident.to_string();
        let field_name_str_lit = LitStr::new(&field_name_str, field_name_ident.span()); // Corrected: Use LitStr for string literals
        let field_ty = &field.ty;
        let field_value_ident = format_ident!("__field_val_{}", field_name_str);

        let (is_option, option_inner_ty_opt) = get_option_inner_type(field_ty);
        let type_for_vec_check = if is_option { option_inner_ty_opt.unwrap() } else { field_ty };
        let (is_vec, vec_inner_ty_opt) = get_vec_inner_type(type_for_vec_check);

        let assignment_code = if is_option {
            let option_inner_ty = option_inner_ty_opt.expect("Option inner type not found");
            if is_vec { // Option<Vec<T>>
                let vec_element_ty = vec_inner_ty_opt.expect("Option<Vec<T>> inner T not found");
                let vec_element_ty_str = get_simple_type_str(vec_element_ty);

                let element_conversion_logic = if is_primitive_type_str(&vec_element_ty_str) {
                    quote! { 
                        let el_for_cast = el.clone();
                        let actual_type_name = el.type_name();
                        el_for_cast.try_cast::<#vec_element_ty>().ok_or_else(|| format!(
                            "Element in Option<Vec> for key '{}' is not of type {}, but of type {}", 
                            #field_name_str_lit, stringify!(#vec_element_ty), actual_type_name
                        ))
                    }
                } else {
                    quote! { 
                        let map_val = el.try_cast::<::rhai::Map>().ok_or_else(|| format!("Expected map for element in Option<Vec> for key '{}', got {}", #field_name_str_lit, el.type_name()))?;
                        #vec_element_ty::from_rhai_map(map_val) 
                    }
                };

                quote! {
                    map.get(#field_name_str_lit).cloned().map(|dyn_val_array| {
                        dyn_val_array.into_array().map_err(|e| e.to_string())?.into_iter().map(|el| {
                            #element_conversion_logic
                        }).collect::<Result<Vec<#vec_element_ty>, String>>()
                    }).transpose()?
                }
            } else { // Option<T>
                let option_inner_ty_str = get_simple_type_str(option_inner_ty);
                let conversion_logic = if is_primitive_type_str(&option_inner_ty_str) {
                    quote! { 
                        let val_for_cast = val.clone();
                        let actual_type_name = val.type_name();
                        val_for_cast.try_cast::<#option_inner_ty>().ok_or_else(|| format!(
                            "Value in Option for key '{}' is not of type {}, but of type {}", 
                            #field_name_str_lit, stringify!(#option_inner_ty), actual_type_name
                        ))
                    }
                } else {
                    quote! { 
                        let map_val = val.try_cast::<::rhai::Map>().ok_or_else(|| format!("Expected map for Option key '{}', got {}", #field_name_str_lit, val.type_name()))?;
                        #option_inner_ty::from_rhai_map(map_val) 
                    }
                };
                quote! {
                    map.get(#field_name_str_lit).cloned().map(|val| { 
                        #conversion_logic
                    }).transpose()?
                }
            }
        } else if is_vec { // Vec<T>
            let vec_inner_ty = vec_inner_ty_opt.expect("Vec inner type not found");
            let vec_inner_ty_str = get_simple_type_str(vec_inner_ty);

            let element_conversion_logic = if is_primitive_type_str(&vec_inner_ty_str) {
                quote! { 
                    let el_for_cast = el.clone();
                    let actual_type_name = el.type_name();
                    el_for_cast.try_cast::<#vec_inner_ty>().ok_or_else(|| format!(
                        "Element in Vec for key '{}' is not of type {}, but of type {}", 
                        #field_name_str_lit, stringify!(#vec_inner_ty), actual_type_name
                    ))
                }
            } else {
                quote! { 
                    let map_val = el.try_cast::<::rhai::Map>().ok_or_else(|| format!("Expected map for element in Vec for key '{}', got {}", #field_name_str_lit, el.type_name()))?;
                    #vec_inner_ty::from_rhai_map(map_val) 
                }
            };

            quote! {{
                let dyn_val_array = map.get(#field_name_str_lit).cloned()
                    .ok_or_else(|| format!("Key '{}' not found for Vec", #field_name_str_lit))?
                    .into_array().map_err(|e| e.to_string())?;
                dyn_val_array.into_iter().map(|el| {
                    #element_conversion_logic
                }).collect::<Result<Vec<#vec_inner_ty>, String>>()?
            }}
        } else { // Direct field T
            let field_ty_str = get_simple_type_str(field_ty);
            let conversion_logic = if is_primitive_type_str(&field_ty_str) {
                quote! { 
                    let dyn_val_for_cast = dyn_val_cloned.clone();
                    let actual_type_name = dyn_val_cloned.type_name();
                    dyn_val_for_cast.try_cast::<#field_ty>().ok_or_else(|| format!(
                        "Value in map for key '{}' is not of type {}, but of type {}", 
                        #field_name_str_lit, stringify!(#field_ty), actual_type_name
                    ))
                }
            } else {
                quote! { 
                    let map_val = dyn_val_cloned.try_cast::<::rhai::Map>().ok_or_else(|| format!("Expected map for key '{}', got {}", #field_name_str_lit, dyn_val_cloned.type_name()))?;
                    #field_ty::from_rhai_map(map_val) 
                }
            };
            quote! {{
                let dyn_val_cloned = map.get(#field_name_str_lit).cloned()
                    .ok_or_else(|| format!("Key '{}' not found", #field_name_str_lit))?;
                #conversion_logic?                
            }}
        };

        field_value_declarations.push(quote! { let #field_value_ident = #assignment_code; });
        struct_field_assignments.push(quote_spanned!(field_name_ident.span()=> #field_name_ident: #field_value_ident));
    }

    let gen = quote! {
        impl #name {
            pub fn from_rhai_map(map: ::rhai::Map) -> Result<Self, String> {
                #(#field_value_declarations)*
                Ok(Self {
                    #(#struct_field_assignments),*
                })
            }
        }
    };

    proc_macro::TokenStream::from(gen)
}

#[proc_macro_derive(ToRhaiMap)]
pub fn derive_to_rhai_map(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let ast = syn::parse_macro_input!(input as syn::DeriveInput);
    let name = &ast.ident;

    let mut field_insertions = Vec::new();

    if let syn::Data::Struct(syn::DataStruct { fields: syn::Fields::Named(fields), .. }) = ast.data {
        for field in fields.named {
            let field_name = field.ident.as_ref().unwrap();
            let field_name_str = field_name.to_string();
            let ty = &field.ty;

            let mut _is_vec_of_custom_struct = false;
            let mut is_direct_custom_struct = false;

            let _field_type_name_for_logic = quote!(#ty).to_string().replace(" ", "");

            if let syn::Type::Path(type_path) = ty {
                if type_path.path.segments.len() == 1 && type_path.path.segments.first().unwrap().ident == "Vec" {
                    if let syn::PathArguments::AngleBracketed(angle_bracketed_args) = &type_path.path.segments.first().unwrap().arguments {
                        if let Some(syn::GenericArgument::Type(_inner_ty_syn @ syn::Type::Path(inner_type_path))) = angle_bracketed_args.args.first() {
                             if let Some(inner_segment) = inner_type_path.path.segments.last() {
                                 let inner_type_str = inner_segment.ident.to_string();
                                 if !["String", "INT", "i64", "FLOAT", "f64", "bool"].contains(&inner_type_str.as_str()) {
                                    _is_vec_of_custom_struct = true;
                                 }
                             }
                         }
                    }
                } else if !["String", "INT", "i64", "FLOAT", "f64", "bool"].contains(&type_path.path.segments.first().unwrap().ident.to_string().as_str()) {
                    // It's not a Vec, and not a primitive, so assume it's a direct custom struct
                    is_direct_custom_struct = true;
                }
            }

            if _is_vec_of_custom_struct {
                field_insertions.push(quote! {
                    {
                        let rhai_array: rhai::Array = self.#field_name.iter().map(|item| {
                            // item is &CustomStruct, to_rhai_map() takes &self
                            rhai::Dynamic::from(item.to_rhai_map())
                        }).collect();
                        map.insert(#field_name_str.into(), rhai::Dynamic::from(rhai_array));
                    }
                });
            // FIX: Logic for direct custom struct field
            } else if is_direct_custom_struct {
                field_insertions.push(quote! {
                    map.insert(#field_name_str.into(), rhai::Dynamic::from(self.#field_name.to_rhai_map()));
                });
            } else {
                // This handles Vec<Primitive> and direct Primitives
                let type_str = quote!(#ty).to_string().replace(" ", "");
                if type_str.starts_with("Vec<") { 
                     field_insertions.push(quote! {
                        {
                            let rhai_array: rhai::Array = self.#field_name.iter().map(|item| {
                                // item is &Primitive, clone it for .into()
                                item.clone().into()
                            }).collect();
                            map.insert(#field_name_str.into(), rhai::Dynamic::from(rhai_array));
                        }
                    });
                } else { 
                    // Direct primitive field
                    field_insertions.push(quote! {
                         map.insert(#field_name_str.into(), self.#field_name.clone().into()); 
                    });
                }
            }
        }
    }

    let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();

    let expanded = quote! {
        impl #impl_generics ToRhaiMap for #name #ty_generics #where_clause {
            fn to_rhai_map(&self) -> rhai::Map {
                let mut map = rhai::Map::new();
                #(#field_insertions)*
                map
            }
        }
    };
    proc_macro::TokenStream::from(expanded)
}

// Rest of the code remains the same