use quote::quote;
use syn::{Data, DeriveInput, Type, parse_str};

use crate::utils::{get_enum_groups, get_struct_groups};

pub fn derive_to_grammar_impl(input: DeriveInput) -> proc_macro2::TokenStream {
    let struct_name = &input.ident; // The name of the type (e.g., "MyStruct")
    let struct_label = struct_name.to_string();

    // Split generics for use in `impl <...> ...`
    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();

    // Generate the body of the `grammar()` function based on data type
    let grammar_body = match &input.data {
        Data::Struct(data) => {
            // Get logical groups (includes bitfield merging logic)
            let groups = get_struct_groups(data);

            let mut main_seq = vec![]; // Tokens for the main rule: [nt("S.f1"), nt("S.f2")]
            let mut extra_rules = vec![]; // Definition logic for fields

            for group in groups {
                // 1. Create a unique Label for this field's rule
                // e.g., "MyStruct.username" or "MyTuple.0"
                let unique_rule_name = format!("{}.{}", struct_label, group.label);
                // 2. Add reference to the Main Sequence
                // The main struct rule will point to this intermediate rule.
                main_seq.push(quote! { aglaea::nt(#unique_rule_name) });

                // If has type override (#[weave(type="...")])
                if let Some(custom_expr) = &group.grammar_override {
                    extra_rules.push(quote! {
                        rules.insert(#unique_rule_name.to_string(), vec![aglaea::exp(vec![aglaea::#custom_expr])]);
                    });
                    continue;
                }

                let base_type = if group.is_bitfield {
                    if group.max_bits <= 8 {
                        quote! { aglaea::U8 }
                    } else if group.max_bits <= 16 {
                        quote! { aglaea::U16be }
                    } else if group.max_bits <= 32 {
                        quote! { aglaea::U32be }
                    } else {
                        quote! { aglaea::U64be }
                    }
                } else {
                    let ty = group.type_origin.as_ref().unwrap();
                    quote! { #ty }
                };

                // Determine target type for grammar generation
                let target_type = if let Some(inner_type_str) = &group.inner_type_annotation {
                    // If has inner type annotation, parse and use it
                    let inner_type: Type = parse_str(inner_type_str)
                        .expect(&format!("Failed to parse inner type: {}", inner_type_str));
                    quote! { #inner_type }
                } else {
                    base_type
                };

                // Register the Intermediate Rule
                extra_rules.push(quote! {
                    rules = rules.extend_grammar(&<#target_type as aglaea::ToGrammar>::grammar_with_context(visited));
                    rules.insert(#unique_rule_name.to_string(), vec![aglaea::exp(vec![
                        aglaea::nt(&<#target_type as aglaea::PangLabel>::label())
                    ])]);
                });
            }

            quote! {
                use aglaea::{exp, nt};
                // Insert Main Rule: "MyStruct" -> [ "MyStruct.field1", "MyStruct.field2" ]
                rules.insert(#struct_label.to_string(), vec![exp(vec![
                    #(#main_seq),*
                ])]);

                // Execute field definitions
                #(#extra_rules)*
            }
        }
        Data::Enum(data) => {
            let enum_groups = get_enum_groups(data);

            let variants_process_logic = enum_groups.into_iter().map(|(variant_name, groups)| {
                // Label for the variant, e.g., "MyEnum.VariantA"
                let variant_label = format!("{}.{}", struct_label, variant_name);

                let mut variant_seq = vec![]; // Sequence of symbols for this variant
                let mut extra_rules = vec![]; // Helper rules (field definitions)

                for group in groups {
                    // Create a unique label for this field/group
                    // e.g., "MyEnum.VariantA.field_name" or "MyEnum.VariantA.flags"
                    let group_rule_name = format!("{}.{}", variant_label, group.label);

                    // Add reference to the Variant Sequence
                    variant_seq.push(quote! { aglaea::nt(#group_rule_name) });

                    // Check for grammar override (#[weave(type="...")])
                    if let Some(custom_expr) = &group.grammar_override {
                        extra_rules.push(quote! {
                            rules.insert(#group_rule_name.to_string(), vec![aglaea::exp(vec![aglaea::#custom_expr])]);
                        });
                        continue;
                    }

                    // Determine the base type (handle Bitfields vs Normal types)
                    let base_type = if group.is_bitfield {
                        if group.max_bits <= 8 {
                            quote! { aglaea::U8 }
                        } else if group.max_bits <= 16 {
                            quote! { aglaea::U16be }
                        } else if group.max_bits <= 32 {
                            quote! { aglaea::U32be }
                        } else {
                            quote! { aglaea::U64be }
                        }
                    } else {
                        let ty = group.type_origin.as_ref().unwrap();
                        quote! { #ty }
                    };

                    // Determine target type for grammar generation
                    let target_type = if let Some(inner_type_str) = &group.inner_type_annotation {
                        // If has inner type annotation, parse and use it
                        let inner_type: Type = parse_str(inner_type_str)
                            .expect(&format!("Failed to parse inner type: {}", inner_type_str));
                        quote! { #inner_type }
                    } else {
                        base_type
                    };

                    // Register the Intermediate Rule for this field
                    // 1. Extend grammar with the target type's sub-grammar
                    // 2. Define the rule: "MyEnum.VariantA.field" -> target_type::label()
                    extra_rules.push(quote! {
                        rules = rules.extend_grammar(&<#target_type as aglaea::ToGrammar>::grammar_with_context(visited));
                        rules.insert(#group_rule_name.to_string(), vec![aglaea::exp(vec![
                            aglaea::nt(&<#target_type as aglaea::PangLabel>::label())
                        ])]);
                    });
                }


                quote! {
                    {
                        // 1. Execute field definitions (side effects & sub-grammars)
                        #(#extra_rules)*

                        // 2. Register the Variant's rule definition
                        // "MyEnum.VariantA" -> [ "MyEnum.VariantA.field1", "MyEnum.VariantA.flags", ... ]
                        rules.insert(#variant_label.to_string(), vec![aglaea::exp(vec![
                            #(#variant_seq),*
                        ])]);

                        // 3. Add this variant as a choice to the parent Enum's list
                        alternatives.push(aglaea::exp(vec![ aglaea::nt(#variant_label) ]));
                    }
                }
            });

            quote! {
                use aglaea::{exp, nt};
                let mut alternatives = vec![];

                // Process all variants
                #(#variants_process_logic)*

                // Insert Main Enum Rule: "MyEnum" -> ( VariantA | VariantB | ... )
                rules.insert(#struct_label.to_string(), alternatives);
            }
        }
        _ => panic!("ToGrammar only supports Struct and Enum"),
    };

    // Wrap everything in the Trait Implementation
    quote! {
        impl #impl_generics aglaea::ToGrammar for #struct_name #ty_generics #where_clause {
            fn grammar_with_context(visited: &mut std::collections::HashSet<String>) -> aglaea::Grammar {
                use aglaea::Grammar;
                let mut rules = Grammar::new();
                let my_label = <Self as aglaea::PangLabel>::label().to_string();

                // Check if already visited to avoid infinite recursion
                if !visited.insert(my_label.clone()) {
                    return rules;
                }

                #grammar_body   // Insert generated rules above
                rules
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use syn::parse_quote;

    #[test]
    fn test_to_grammar_middleware_naming() {
        let input: DeriveInput = parse_quote! {
            struct Header {
                ver: u8,
                len: u16
            }
        };
        let output = derive_to_grammar_impl(input);
        let s = output.to_string();

        // As test above
        assert!(s.contains("aglaea :: nt (\"Header.ver\")"));
        assert!(s.contains("aglaea :: nt (\"Header.len\")"));

        // rules.extend_grammar(&u8::grammar_with_context(visited));
        assert!(s.contains("rules . extend_grammar (& < u8 as aglaea :: ToGrammar > :: grammar_with_context (visited))"));

        // "Header.ver" => [exp(["u8"])]
        assert!(s.contains("vec ! [aglaea :: nt (& < u8 as aglaea :: PangLabel > :: label ())]"));

        // rules.extend_grammar(&u16::grammar_with_context(visited));
        assert!(
            s.contains("rules . extend_grammar (& < u16 as aglaea :: ToGrammar > :: grammar_with_context (visited))")
        );

        // "Header.len" => [exp(["u16"])]
        assert!(s.contains("vec ! [aglaea :: nt (& < u16 as aglaea :: PangLabel > :: label ())]"));
    }
}
