use crate::spec::tree_enum_spec::{EnumRuleChildKind, TreeEnumBuildSpec};
use convert_case::{Case, Casing};
use proc_macro2::TokenStream;
use quote::{format_ident, quote};

pub fn make_enum_ast_node_impl(enum_spec: &TreeEnumBuildSpec) -> TokenStream {
    let type_ident = format_ident!("{}", enum_spec.build());
    let match_arms = enum_spec
        .rules()
        .map(|rule| {
            let rule_ident = format_ident!("{}", rule.rule());
            match rule.child() {
                Some(child) => match child.kind() {
                    EnumRuleChildKind::Node(node_child) => {
                        let child_ident =
                            format_ident!("{}", node_child.node_kind().to_case(Case::Snake));
                        quote! {
                            #type_ident::#rule_ident(#child_ident) => vec![
                                #child_ident
                            ]
                        }
                    }
                    _ => quote! {
                        #type_ident::#rule_ident(_) => vec![]
                    },
                },
                None => {
                    quote! {
                        #type_ident::#rule_ident => vec![]
                    }
                }
            }
        })
        .collect::<Vec<_>>();

    let mut_match_arms = enum_spec
        .rules()
        .map(|rule| {
            let rule_ident = format_ident!("{}", rule.rule());
            match rule.child() {
                Some(child) => match child.kind() {
                    EnumRuleChildKind::Node(node_child) => {
                        let child_ident =
                            format_ident!("{}", node_child.node_kind().to_case(Case::Snake));
                        quote! {
                            #type_ident::#rule_ident(#child_ident) => {
                                f(#child_ident as &'a mut dyn AstNode<'a>)
                            }
                        }
                    }
                    _ => quote! {
                        #type_ident::#rule_ident(_) => {}
                    },
                },
                None => {
                    quote! {
                        #type_ident::#rule_ident => {}
                    }
                }
            }
        })
        .collect::<Vec<_>>();

    quote! {
        impl<'a> AstNode<'a> for #type_ident {
            fn children(&'a self) -> Vec<&'a dyn AstNode<'a>> {
                match self {
                    #(#match_arms,)*
                }
            }

            fn for_each_child_mut(&'a mut self, mut f: &mut dyn FnMut(&'a mut dyn AstNode<'a>)) {
                match self {
                    #(#mut_match_arms,)*
                }
            }

            fn as_node_ref(&'a self) -> AstNodeRef<'a> {
                AstNodeRef::#type_ident(&self)
            }

            fn as_node_ref_mut(&'a mut self) -> AstNodeRefMut<'a> {
                AstNodeRefMut::#type_ident(self)
            }
        }
    }
}