use quote::quote;
use syn::{Data, DeriveInput, Fields};

use crate::utils::{get_struct_groups};

pub fn derive_weaver_impl(input: DeriveInput) -> proc_macro2::TokenStream {
    let struct_name = &input.ident;
    let struct_label = struct_name.to_string();
    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();

    let body = match &input.data {
        Data::Struct(data) => {
            let groups = get_struct_groups(data);

            let child_nodes = groups.iter().map(|group| {
                let node_name = format!("{}.{}", struct_label, group.label);

                let generation_logic = if group.is_bitfield {
                    // --- Process Bitfield ---
                    if group.max_bits <= 8 {
                        quote! { <u8 as fugue_fuzz::Weaver>::weave(s) }
                    } else if group.max_bits <= 16 {
                        quote! { <u16 as fugue_fuzz::Weaver>::weave(s) }
                    } else if group.max_bits <= 32 {
                        quote! { <u32 as fugue_fuzz::Weaver>::weave(s) }
                    } else {
                        quote! { <u64 as fugue_fuzz::Weaver>::weave(s) }
                    }
                } else {
                    let ty = group.type_origin.as_ref().unwrap();
                    let _attrs = &data.fields.iter().nth(group.members[0]).unwrap().attrs;

                    // Check #[weave(type="...")] override
                    if let Some(custom_expr) = &group.grammar_override {
                        quote! { <#custom_expr as fugue_fuzz::Weaver>::weave(s) }
                    } else {
                        // Check #[weave(inner="...")] override
                        let target_type = if let Some(inner_str) = &group.inner_type_annotation {
                            let inner_type: syn::Type = syn::parse_str(inner_str)
                                .expect(&format!("Failed to parse inner type: {}", inner_str));
                            quote! { #inner_type }
                        } else {
                            quote! { #ty }
                        };

                        // TODO: Check #[weave(range="...")]
                        // let range = get_range_from_attrs(attrs);
                        // if let Some(r) = range { ... }

                        quote! { <#target_type as fugue_fuzz::Weaver>::weave(s) }
                    }
                };

                quote! {
                    aglaea::new_node(aglaea::ntc(#node_name, Some(vec![#generation_logic])))
                }
            });

            quote! {
                let children = vec![
                    #(#child_nodes),*
                ];
                aglaea::new_node(aglaea::ntc(#struct_label, Some(children)))
            }
        }
        Data::Enum(data) => {
            let variants_count = data.variants.len();
            if variants_count == 0 {
                quote! { panic!("Cannot weave an empty enum!"); }
            } else {
                let enum_groups = crate::utils::get_enum_groups(data);

                let match_arms = enum_groups.iter().enumerate().map(|(idx, (variant_name, groups))| {
                    let variant_label = format!("{}.{}", struct_label, variant_name);

                    let variant = &data.variants[idx];

                    let child_nodes = groups.iter().map(|group| {
                        let node_name = format!("{}.{}", variant_label, group.label);

                        let generation_logic = if group.is_bitfield {
                            // Process Bitfield (Same as Struct)
                            if group.max_bits <= 8 {
                                quote! { <u8 as fugue_fuzz::Weaver>::weave(s) }
                            } else if group.max_bits <= 16 {
                                quote! { <u16 as fugue_fuzz::Weaver>::weave(s) }
                            } else if group.max_bits <= 32 {
                                quote! { <u32 as fugue_fuzz::Weaver>::weave(s) }
                            } else {
                                quote! { <u64 as fugue_fuzz::Weaver>::weave(s) }
                            }
                        } else {
                            let ty = group.type_origin.as_ref().unwrap();
                            let member_idx = group.members[0];
                            let _attrs = match &variant.fields {
                                Fields::Named(f) => &f.named[member_idx].attrs,
                                Fields::Unnamed(f) => &f.unnamed[member_idx].attrs,
                                Fields::Unit => &vec![],
                            };

                            if let Some(custom_expr) = &group.grammar_override {
                                quote! { <#custom_expr as fugue_fuzz::Weaver>::weave(s) }
                            } else {
                                // Check #[weave(inner="...")] override
                                let target_type = if let Some(inner_str) = &group.inner_type_annotation {
                                    let inner_type: syn::Type = syn::parse_str(inner_str)
                                        .expect(&format!("Failed to parse inner type: {}", inner_str));
                                    quote! { #inner_type }
                                } else {
                                    quote! { #ty }
                                };
                                quote! { <#target_type as fugue_fuzz::Weaver>::weave(s) }
                            }
                        };

                        quote! {
                            aglaea::new_node(aglaea::ntc(#node_name, Some(vec![#generation_logic])))
                        }
                    });

                    quote! {
                        #idx => {
                            let children = vec![ #(#child_nodes),* ];
                            let variant_node = aglaea::new_node(aglaea::ntc(#variant_label, Some(children)));
                            aglaea::new_node(aglaea::ntc(#struct_label, Some(vec![variant_node])))
                        }
                    }
                });

                quote! {
                    let choice = s.rand_mut().below_or_zero(#variants_count);
                    match choice {
                        #(#match_arms),*
                        _ => unreachable!(),
                    }
                }
            }
        }
        _ => panic!("Weaver only supports Struct and Enum"),
    };

    quote! {
        impl #impl_generics fugue_fuzz::Weaver for #struct_name #ty_generics #where_clause {
            fn weave<S>(s: &mut S) -> std::sync::Arc<aglaea::DerivationTree>
            where S: libafl::state::HasRand
            {
                #body
            }
        }
    }
}
