db/heromodels-derive/src/lib.rs

use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{Data, DeriveInput, Fields, parse_macro_input};

/// Convert a string to snake_case
fn to_snake_case(s: &str) -> String {
    let mut result = String::new();
    for (i, c) in s.char_indices() {
        if i > 0 && c.is_uppercase() {
            result.push('_');
        }
        result.push(c.to_lowercase().next().unwrap());
    }
    result
}

/// Convert a string to PascalCase
// fn to_pascal_case(s: &str) -> String {
//     let mut result = String::new();
//     let mut capitalize_next = true;
//
//     for c in s.chars() {
//         if c == '_' {
//             capitalize_next = true;
//         } else if capitalize_next {
//             result.push(c.to_uppercase().next().unwrap());
//             capitalize_next = false;
//         } else {
//             result.push(c);
//         }
//     }
//
//     result
// }

/// Implements the Model trait and generates Index trait implementations for fields marked with #[index].
#[proc_macro_attribute]
pub fn model(_attr: TokenStream, item: TokenStream) -> TokenStream {
    // Parse the input tokens into a syntax tree
    let mut input = parse_macro_input!(item as DeriveInput);

    // Extract struct name
    let struct_name = &input.ident;

    // Convert struct name to snake_case for db_prefix
    let name_str = struct_name.to_string();
    let db_prefix = to_snake_case(&name_str);

    // Extract fields with #[index] attribute
    let mut indexed_fields = Vec::new();
    let mut custom_index_names = std::collections::HashMap::new();

    if let Data::Struct(ref mut data_struct) = input.data {
        if let Fields::Named(ref mut fields_named) = data_struct.fields {
            for field in &mut fields_named.named {
                let mut attr_idx = None;
                for (i, attr) in field.attrs.iter().enumerate() {
                    if attr.path().is_ident("index") {
                        attr_idx = Some(i);
                        if let Some(ref field_name) = field.ident {
                            // Check if the attribute has parameters
                            let mut custom_name = None;

                            // Parse attribute arguments if any
                            let meta = attr.meta.clone();
                            if let syn::Meta::List(list) = meta {
                                if let Ok(nested) = list.parse_args_with(syn::punctuated::Punctuated::<syn::Meta, syn::Token![,]>::parse_terminated) {
                                        for meta in nested {
                                            if let syn::Meta::NameValue(name_value) = meta {
                                                if name_value.path.is_ident("name") {
                                                    if let syn::Expr::Lit(syn::ExprLit { lit: syn::Lit::Str(lit_str), .. }) = name_value.value {
                                                        custom_name = Some(lit_str.value());
                                                    }
                                                }
                                            }
                                        }
                                    }
                            }

                            indexed_fields.push((field_name.clone(), field.ty.clone()));

                            if let Some(name) = custom_name {
                                custom_index_names.insert(field_name.to_string(), name);
                            }
                        }
                    }
                }

                if let Some(idx) = attr_idx {
                    field.attrs.remove(idx);
                }
            }
        }
    }

    // Generate Model trait implementation
    let db_keys_impl = if indexed_fields.is_empty() {
        quote! {
            fn db_keys(&self) -> Vec<heromodels_core::IndexKey> {
                Vec::new()
            }
        }
    } else {
        let field_keys = indexed_fields.iter().map(|(field_name, _)| {
            let name_str = custom_index_names
                .get(&field_name.to_string())
                .cloned()
                .unwrap_or(field_name.to_string());
            quote! {
                heromodels_core::IndexKey {
                    name: #name_str,
                    value: self.#field_name.to_string(),
                }
            }
        });

        quote! {
            fn db_keys(&self) -> Vec<heromodels_core::IndexKey> {
                vec![
                    #(#field_keys),*
                ]
            }
        }
    };

    let model_impl = quote! {
        impl heromodels_core::Model for #struct_name {
            fn db_prefix() -> &'static str {
                #db_prefix
            }

            fn get_id(&self) -> u32 {
                self.base_data.id
            }

            fn base_data_mut(&mut self) -> &mut heromodels_core::BaseModelData {
                &mut self.base_data
            }

            #db_keys_impl
        }
    };

    // Generate Index trait implementations
    let mut index_impls = proc_macro2::TokenStream::new();

    for (field_name, field_type) in &indexed_fields {
        let name_str = field_name.to_string();

        // Get custom index name if specified, otherwise use field name
        let index_key = match custom_index_names.get(&name_str) {
            Some(custom_name) => custom_name.clone(),
            None => name_str.clone(),
        };

        // Convert field name to PascalCase for struct name
        // let struct_name_str = to_pascal_case(&name_str);
        // let index_struct_name = format_ident!("{}", struct_name_str);
        let index_struct_name = format_ident!("{}", &name_str);

        // Default to str for key type
        let index_impl = quote! {
            pub struct #index_struct_name;

            impl heromodels_core::Index for #index_struct_name {
                type Model = super::#struct_name;
                type Key = #field_type;

                fn key() -> &'static str {
                    #index_key
                }
            }
        };

        index_impls.extend(index_impl);
    }

    if !index_impls.is_empty() {
        let index_mod_name = format_ident!("{}_index", db_prefix);
        index_impls = quote! {
            pub mod #index_mod_name {
                #index_impls
            }
        }
    }

    // Combine the original struct with the generated implementations
    let expanded = quote! {
        #input

        #model_impl

        #index_impls
    };

    // Return the generated code
    expanded.into()
}