use std::collections::HashSet;

use inflector::Inflector;
use proc_macro2::{Span, TokenStream};
use quote::{quote, ToTokens};
use syn::{
    parse_quote, Arm, Attribute, Expr, Field, Fields, File, GenericArgument, Ident, Item, Lit,
    LitInt, Path, PathArguments, Stmt, TraitItem, Type,
};

pub fn generate(crate_name: &Ident, node_types: &[&Item]) -> File {
    let mut output = File {
        shebang: None,
        attrs: Vec::new(),
        items: Vec::new(),
    };
    let mut all_types = all_field_types(node_types).into_iter().collect::<Vec<_>>();
    all_types.sort_by_cached_key(|v| v.method_name());

    let mut typedefs = HashSet::new();

    for node_type in node_types {
        match node_type {
            Item::Enum(data) => {
                typedefs.insert(FieldType::Normal(data.ident.to_string()));
            }
            Item::Struct(data) => {
                typedefs.insert(FieldType::Normal(data.ident.to_string()));
            }
            _ => {}
        }
    }

    let field_only_types = {
        let mut all = all_types.clone();

        all.retain(|ty| !typedefs.contains(ty));

        all
    };

    output.attrs.push(parse_quote!(
        //! This file is generated by `tools/generate-code`. DO NOT MODIFY.
    ));
    output.attrs.push(parse_quote!(
        #![allow(unused_variables)]
    ));
    output.attrs.push(parse_quote!(
        #![allow(clippy::all)]
    ));

    output.items.push(parse_quote!(
        use #crate_name::*;
    ));

    output.items.push(parse_quote!(
        pub use ::swc_visit::All;
    ));

    for &kind in [TraitKind::Visit, TraitKind::VisitMut, TraitKind::Fold].iter() {
        for &variant in [Variant::Normal, Variant::AstPath].iter() {
            let g = Generator { kind, variant };

            output.items.extend(g.declare_visit_trait(&all_types));

            output.items.extend(g.declare_visit_with_trait());

            output
                .items
                .extend(g.implement_visit_with_for_node_types(node_types));

            output
                .items
                .extend(g.implement_visit_with_for_non_node_types(&field_only_types));

            output
                .items
                .extend(g.implement_visit_with_for_generic_types());
        }
    }

    output.items.push(parse_quote!(
        #[cfg(any(docsrs, feature = "path"))]
        pub type AstKindPath = swc_visit::AstKindPath<AstParentKind>;
    ));
    output.items.push(parse_quote!(
        #[cfg(any(docsrs, feature = "path"))]
        pub type AstNodePath<'ast> = swc_visit::AstNodePath<AstParentNodeRef<'ast>>;
    ));
    output.items.extend(define_fields(crate_name, node_types));

    output
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
enum FieldType {
    Normal(String),
    Generic(String, Box<FieldType>),
}

impl ToTokens for FieldType {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        match self {
            FieldType::Normal(name) => {
                let parsed: Path = syn::parse_str(name).expect("failed to parse path");
                parsed.to_tokens(tokens);
            }
            FieldType::Generic(name, ty) => {
                let name = Ident::new(name, Span::call_site());
                let ty = &**ty;
                quote!(#name<#ty>).to_tokens(tokens);
            }
        }
    }
}

impl FieldType {
    pub fn method_name(&self) -> String {
        match self {
            FieldType::Normal(name) => name.split("::").last().unwrap().to_snake_case(),
            FieldType::Generic(name, ty) => match &**name {
                "Option" => format!("opt_{}", ty.method_name()),
                "Vec" => {
                    // Vec<Option<Foo>> => opt_vec_foo
                    match &**ty {
                        FieldType::Generic(name, ty) if name == "Option" => {
                            return format!("opt_vec_{}s", ty.method_name())
                        }
                        _ => {}
                    }
                    format!("{}s", ty.method_name())
                }
                "Box" => ty.method_name(),
                _ => todo!("method_name for generic type: {}", name),
            },
        }
    }
}

fn all_field_types(node_types: &[&Item]) -> HashSet<FieldType> {
    let mut all_types = HashSet::new();

    for ty in node_types {
        let type_name = match ty {
            Item::Enum(data) => data.ident.to_string(),
            Item::Struct(data) => data.ident.to_string(),
            _ => continue,
        };
        all_types.insert(FieldType::Normal(type_name));

        match ty {
            Item::Enum(data) => {
                for variant in &data.variants {
                    for field in &variant.fields {
                        let ty = &field.ty;
                        all_types.extend(all_types_in_ty(ty));
                    }
                }
            }
            Item::Struct(data) => {
                for field in &data.fields {
                    let ty = &field.ty;
                    all_types.extend(all_types_in_ty(ty));
                }
            }
            _ => continue,
        }
    }

    all_types
}

fn to_field_ty(ty: &Type) -> Option<FieldType> {
    if let Some(ty) = extract_vec(ty) {
        return to_field_ty(ty).map(|ty| FieldType::Generic("Vec".into(), Box::new(ty)));
    }

    if let Some(ty) = extract_generic("Box", ty) {
        return to_field_ty(ty).map(|ty| FieldType::Generic("Box".into(), Box::new(ty)));
    }

    if let Some(ty) = extract_generic("Option", ty) {
        return to_field_ty(ty).map(|ty| FieldType::Generic("Option".into(), Box::new(ty)));
    }

    match ty {
        Type::Path(p) => {
            let last = p.path.segments.last().unwrap();

            if last.arguments.is_empty() {
                let i = &last.ident;

                if i == "bool"
                    || i == "char"
                    || i == "f32"
                    || i == "f64"
                    || i == "i8"
                    || i == "i16"
                    || i == "i32"
                    || i == "i64"
                    || i == "i128"
                    || i == "isize"
                    || i == "str"
                    || i == "u8"
                    || i == "u16"
                    || i == "u32"
                    || i == "u64"
                    || i == "u128"
                    || i == "usize"
                {
                    return None;
                }

                return Some(FieldType::Normal(quote!(#p).to_string()));
            }

            todo!("to_field_ty: {:?}", ty)
        }
        _ => todo!("to_field_ty"),
    }
}

fn all_types_in_ty(ty: &Type) -> Vec<FieldType> {
    if let Some(ty) = extract_vec(ty) {
        let mut types = all_types_in_ty(ty);
        types.extend(to_field_ty(ty).map(|ty| FieldType::Generic("Vec".into(), Box::new(ty))));
        return types;
    }

    if let Some(ty) = extract_generic("Box", ty) {
        let mut types = all_types_in_ty(ty);
        types.extend(to_field_ty(ty).map(|ty| FieldType::Generic("Box".into(), Box::new(ty))));
        return types;
    }

    if let Some(ty) = extract_generic("Option", ty) {
        let mut types = all_types_in_ty(ty);
        types.extend(to_field_ty(ty).map(|ty| FieldType::Generic("Option".into(), Box::new(ty))));
        return types;
    }

    to_field_ty(ty).into_iter().collect()
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TraitKind {
    Visit,
    VisitMut,
    Fold,
}

impl TraitKind {
    pub fn method_prefix(self) -> &'static str {
        match self {
            TraitKind::Visit => "visit",
            TraitKind::VisitMut => "visit_mut",
            TraitKind::Fold => "fold",
        }
    }

    pub fn trait_prefix(self) -> &'static str {
        match self {
            TraitKind::Visit => "Visit",
            TraitKind::VisitMut => "VisitMut",
            TraitKind::Fold => "Fold",
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Variant {
    Normal,
    AstPath,
}

impl Variant {
    pub fn method_suffix(self, is_visitor_method: bool) -> &'static str {
        if self == Variant::Normal || is_visitor_method {
            ""
        } else {
            "_ast_path"
        }
    }
}

struct Generator {
    kind: TraitKind,
    variant: Variant,
}

impl Generator {
    fn should_skip(&self, ty: &Type) -> bool {
        if let Some(ty) = extract_generic("Box", ty) {
            return self.should_skip(ty);
        }

        if let Some(ty) = extract_generic("Vec", ty) {
            return self.should_skip(ty);
        }

        if let Some(ty) = extract_generic("Option", ty) {
            return self.should_skip(ty);
        }

        let ty = to_field_ty(ty);
        match ty {
            Some(..) => {}
            None => return true,
        }
        false
    }

    fn method_lifetime(&self) -> TokenStream {
        match self.kind {
            TraitKind::Visit => match self.variant {
                Variant::Normal => quote!(),
                Variant::AstPath => quote!(<'ast: 'r, 'r>),
            },
            TraitKind::VisitMut => quote!(),
            TraitKind::Fold => quote!(),
        }
    }

    fn parameter_type_token(&self, ty: TokenStream) -> TokenStream {
        match self.kind {
            TraitKind::Visit => match self.variant {
                Variant::Normal => quote!(&#ty),
                Variant::AstPath => quote!(&'ast #ty),
            },
            TraitKind::VisitMut => quote!(&mut #ty),
            TraitKind::Fold => ty,
        }
    }

    /// This includes `->`
    fn return_type_token(&self, ty: TokenStream) -> TokenStream {
        match self.kind {
            TraitKind::Visit => quote!(),
            TraitKind::VisitMut => quote!(),
            TraitKind::Fold => quote!(-> #ty),
        }
    }

    fn arg_extra_token(&self) -> TokenStream {
        match self.variant {
            Variant::Normal => quote!(),
            Variant::AstPath => quote!(, __ast_path),
        }
    }

    fn param_extra_token(&self) -> TokenStream {
        match self.variant {
            Variant::Normal => quote!(),
            Variant::AstPath => match self.kind {
                TraitKind::Visit => {
                    quote!(, __ast_path: &mut AstNodePath<'r>)
                }
                TraitKind::VisitMut | TraitKind::Fold => quote!(, __ast_path: &mut AstKindPath),
            },
        }
    }

    fn trait_name(&self, with: bool) -> Ident {
        let name = self.kind.trait_prefix();

        let name = if with {
            format!("{}With", name)
        } else {
            name.to_string()
        };

        match self.variant {
            Variant::Normal => Ident::new(&name, Span::call_site()),
            Variant::AstPath => Ident::new(&format!("{}AstPath", name), Span::call_site()),
        }
    }

    fn base_trait_attrs(&self) -> Vec<Attribute> {
        let mut attrs = Vec::new();

        if self.variant == Variant::AstPath {
            attrs.push(parse_quote!(#[cfg(any(docsrs, feature = "path"))]));
            attrs.push(parse_quote!(#[cfg_attr(docsrs, doc(cfg(feature = "path")))]));
        }

        attrs
    }

    fn declare_visit_trait(&self, all_types: &[FieldType]) -> Vec<Item> {
        let mut items = Vec::<Item>::new();
        let lifetime = self.method_lifetime();
        let ast_path_arg = self.arg_extra_token();
        let ast_path_params = self.param_extra_token();
        let with_trait_name = self.trait_name(true);
        let trait_name = self.trait_name(false);
        let attrs = self.base_trait_attrs();
        let mut trait_methods = Vec::<TraitItem>::new();
        let mut either_impl_methods = Vec::<TraitItem>::new();
        let mut optional_impl_methods = Vec::<TraitItem>::new();
        let mut ptr_impl_methods = Vec::<TraitItem>::new();

        for ty in all_types {
            if let FieldType::Generic(name, ..) = &ty {
                if name == "Box" {
                    continue;
                }
            }
            let type_name = quote!(#ty);
            let return_type = self.return_type_token(quote!(#type_name));
            let node_type = self.node_type_for_visitor_method(ty);
            let type_param = self.parameter_type_token(quote!(#node_type));

            let visit_method_name = Ident::new(
                &format!(
                    "{}_{}{}",
                    self.kind.method_prefix(),
                    ty.method_name(),
                    self.variant.method_suffix(true)
                ),
                Span::call_site(),
            );
            let visit_with_children_name = Ident::new(
                &format!(
                    "{}_children_with{}",
                    self.kind.method_prefix(),
                    self.variant.method_suffix(false)
                ),
                Span::call_site(),
            );

            let recurse_doc = "If you want to recurse, you need to call it manually.";

            let method_doc = doc(&format!(
                "Visit a node of type `{}`.\n\nBy default, this method calls \
                 [`{type_name}::{visit_with_children_name}`]. {recurse_doc}",
                type_name
            ));

            trait_methods.push(parse_quote!(
                #method_doc
                #[inline]
                fn #visit_method_name #lifetime (&mut self, node: #type_param #ast_path_params) #return_type {
                    <#node_type as #with_trait_name<Self>>::#visit_with_children_name(node, self #ast_path_arg)
                }
            ));

            either_impl_methods.push(parse_quote!(
                #[inline]
                fn #visit_method_name #lifetime (&mut self, node: #type_param #ast_path_params) #return_type {
                    match self {
                        swc_visit::Either::Left(visitor) => {
                            #trait_name::#visit_method_name(visitor, node #ast_path_arg)
                        }
                        swc_visit::Either::Right(visitor) => {
                            #trait_name::#visit_method_name(visitor, node #ast_path_arg)
                        }
                    }
                }
            ));

            let else_block = if self.kind == TraitKind::Fold {
                quote!(node)
            } else {
                quote!()
            };
            optional_impl_methods.push(parse_quote!(
                #[inline]
                fn #visit_method_name #lifetime (&mut self, node: #type_param #ast_path_params) #return_type {
                    if self.enabled {
                        <V as #trait_name>::#visit_method_name(&mut self.visitor, node #ast_path_arg)
                    } else {
                        #else_block
                    }
                }
            ));

            ptr_impl_methods.push(parse_quote!(
                #[inline]
                fn #visit_method_name #lifetime (&mut self, node: #type_param #ast_path_params) #return_type {
                    <V as #trait_name>::#visit_method_name(&mut **self, node #ast_path_arg)
                }
            ));
        }

        items.push(parse_quote! {
            /// A visitor trait for traversing the AST.
            #(#attrs)*
            pub trait #trait_name {
                #(#trait_methods)*
            }
        });

        // &mut V
        items.push(parse_quote! {
            #(#attrs)*
            impl<V> #trait_name for &mut V where V: ?Sized + #trait_name {
                #(#ptr_impl_methods)*
            }
        });

        // Box<V>
        items.push(parse_quote! {
            #(#attrs)*
            impl<V> #trait_name for Box<V> where V: ?Sized + #trait_name {
                #(#ptr_impl_methods)*
            }
        });

        // ::swc_visit::Either<A, B>

        items.push(parse_quote! {
            #(#attrs)*
            impl<A, B> #trait_name for ::swc_visit::Either<A, B>
            where
                A: #trait_name,
                B: #trait_name,
            {
                #(#either_impl_methods)*
            }
        });

        // ::swc_visit::Optional<V>

        items.push(parse_quote! {
            #(#attrs)*
            impl<V> #trait_name for ::swc_visit::Optional<V>
            where
                V: #trait_name,
            {
                #(#optional_impl_methods)*
            }
        });

        items
    }

    fn declare_visit_with_trait(&self) -> Vec<Item> {
        let visitor_trait_name = self.trait_name(false);
        let trait_name = self.trait_name(true);
        let attrs = self.base_trait_attrs();
        let mut visit_with_trait_methods: Vec<TraitItem> = Vec::new();

        {
            let lifetime = self.method_lifetime();
            let ast_path_extra = self.param_extra_token();
            let return_type = self.return_type_token(quote!(Self));
            let receiver = self.parameter_type_token(quote!(self));

            let visit_with_name = Ident::new(
                &format!(
                    "{}_with{}",
                    self.kind.method_prefix(),
                    self.variant.method_suffix(false)
                ),
                Span::call_site(),
            );
            let visit_with_children_name = Ident::new(
                &format!(
                    "{}_children_with{}",
                    self.kind.method_prefix(),
                    self.variant.method_suffix(false)
                ),
                Span::call_site(),
            );

            visit_with_trait_methods.push(parse_quote!(
                /// Calls a visitor method (visitor.fold_xxx) with self.
                fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_extra) #return_type;
            ));

            visit_with_trait_methods.push(parse_quote!(
            /// Visit children nodes of `self`` with `visitor`.
            fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_extra) #return_type;
        ));
        }

        let mut items: Vec<Item> = Vec::new();
        items.push(parse_quote!(
            /// A trait implemented for types that can be visited using a visitor.
            #(#attrs)*
            pub trait #trait_name<V: ?Sized + #visitor_trait_name> {
                #(#visit_with_trait_methods)*
            }
        ));

        items
    }

    fn implement_visit_with_for_node_types(&self, node_types: &[&Item]) -> Vec<Item> {
        let visitor_trait_name = self.trait_name(false);
        let trait_name = self.trait_name(true);
        let attrs = self.base_trait_attrs();

        let mut items: Vec<Item> = Vec::new();

        for node_type in node_types {
            let type_name = match node_type {
                Item::Enum(data) => data.ident.clone(),
                Item::Struct(data) => data.ident.clone(),
                _ => continue,
            };

            let lifetime = self.method_lifetime();
            let ast_path_arg = self.arg_extra_token();
            let ast_path_param = self.param_extra_token();
            let return_type = self.return_type_token(quote!(Self));

            let receiver = self.parameter_type_token(quote!(self));
            let visit_with_name = Ident::new(
                &format!(
                    "{}_with{}",
                    self.kind.method_prefix(),
                    self.variant.method_suffix(false)
                ),
                Span::call_site(),
            );
            let visit_with_children_name = Ident::new(
                &format!(
                    "{}_children_with{}",
                    self.kind.method_prefix(),
                    self.variant.method_suffix(false)
                ),
                Span::call_site(),
            );

            let visit_method_name = Ident::new(
                &format!(
                    "{}_{}{}",
                    self.kind.method_prefix(),
                    type_name.to_string().to_snake_case(),
                    self.variant.method_suffix(true)
                ),
                Span::call_site(),
            );

            let visit_with_doc = doc(&format!(
                "Calls [{visitor_trait_name}`::{}`] with `self`.",
                visit_method_name
            ));

            let default_body: Expr = match node_type {
                Item::Enum(data) => {
                    let name = &data.ident;
                    let mut match_arms = Vec::new();

                    for v in &data.variants {
                        let variant_name = &v.ident;

                        match_arms.push(self.default_visit_body(
                            quote!(#name::#variant_name),
                            name,
                            Some(variant_name),
                            &v.fields,
                        ));
                    }

                    parse_quote!(match self { #(#match_arms)* })
                }
                Item::Struct(data) => {
                    let name = &data.ident;

                    let arm = self.default_visit_body(quote!(#name), name, None, &data.fields);

                    parse_quote!(match self { #arm })
                }
                _ => continue,
            };

            items.push(parse_quote!(
                #(#attrs)*
                impl<V: ?Sized + #visitor_trait_name> #trait_name<V> for #type_name {
                    #visit_with_doc
                    fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
                        <V as #visitor_trait_name>::#visit_method_name(visitor, self #ast_path_arg)
                    }

                    fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
                        #default_body
                    }
                }
            ));
        }

        items
    }

    fn default_visit_body(
        &self,
        path: TokenStream,
        type_name: &Ident,
        enum_variant_name: Option<&Ident>,
        fields: &Fields,
    ) -> Arm {
        let ast_path_arg = match self.variant {
            Variant::Normal => quote!(),
            Variant::AstPath => quote!(, &mut *__ast_path),
        };

        let with_visitor_trait_name = self.trait_name(true);
        let visit_with_name = Ident::new(
            &format!(
                "{}_with{}",
                self.kind.method_prefix(),
                self.variant.method_suffix(false)
            ),
            Span::call_site(),
        );

        let fields_enum_name = Ident::new(&format!("{type_name}Field"), Span::call_site());

        let enum_ast_path = match enum_variant_name {
            Some(variant_name) if self.variant == Variant::AstPath => {
                let field_variant = Ident::new(
                    &variant_name.to_string().to_pascal_case(),
                    Span::call_site(),
                );

                match self.kind {
                    TraitKind::Visit => Some(quote!(
                        let mut __ast_path = __ast_path
                            .with_guard(
                                AstParentNodeRef::#type_name(self, self::fields::#fields_enum_name::#field_variant),
                            );
                    )),
                    _ => Some(quote!(
                        let mut __ast_path = __ast_path
                            .with_guard(
                                AstParentKind::#type_name(self::fields::#fields_enum_name::#field_variant),
                            );
                    )),
                }
            }
            _ => None,
        };

        match fields {
            Fields::Named(n) => {
                let mut stmts: Vec<Stmt> = Vec::new();
                let mut bindings = Vec::new();
                let mut reconstruct = match self.kind {
                    TraitKind::Visit | TraitKind::VisitMut => None,
                    TraitKind::Fold => Some(Vec::<TokenStream>::new()),
                };

                for field in &n.named {
                    let field_name = field.ident.as_ref().unwrap();
                    let ty = &field.ty;

                    bindings.push(field_name.clone());

                    let field_variant =
                        Ident::new(&field_name.to_string().to_pascal_case(), Span::call_site());

                    let mut ast_path_guard_expr: Option<Stmt> = None;

                    if self.variant == Variant::AstPath && !self.should_skip(ty) {
                        let mut kind = quote!(self::fields::#fields_enum_name::#field_variant);

                        if extract_vec(extract_generic("Option", ty).unwrap_or(ty)).is_some() {
                            kind = quote!(#kind(usize::MAX));
                        }

                        match self.kind {
                            TraitKind::Visit => {
                                ast_path_guard_expr = Some(parse_quote!(
                                    let mut __ast_path = __ast_path
                                        .with_guard(
                                            AstParentNodeRef::#type_name(self, #kind),
                                        );
                                ));
                            }
                            _ => {
                                ast_path_guard_expr = Some(parse_quote!(
                                    let mut __ast_path = __ast_path
                                        .with_guard(
                                            AstParentKind::#type_name(#kind),
                                        );
                                ));
                            }
                        }
                    }

                    if let Some(reconstructor) = &mut reconstruct {
                        if !self.should_skip(ty) {
                            stmts.push(parse_quote!(
                                let #field_name = {
                                    #ast_path_guard_expr
                                    <#ty as #with_visitor_trait_name<V>>::#visit_with_name(#field_name, visitor #ast_path_arg)
                                };
                            ));
                        }

                        reconstructor.push(parse_quote!(#field_name));
                    } else if !self.should_skip(ty) {
                        stmts.push(parse_quote!(
                            {
                                #ast_path_guard_expr
                                <#ty as #with_visitor_trait_name<V>>::#visit_with_name(#field_name, visitor #ast_path_arg)
                            };
                        ));
                    }
                }

                match self.kind {
                    TraitKind::Visit | TraitKind::VisitMut => {
                        parse_quote!(#path { #(#bindings),* } => {
                            #enum_ast_path;

                            #(#stmts)*
                        })
                    }
                    TraitKind::Fold => {
                        let reconstruct = reconstruct.unwrap();

                        parse_quote!(#path { #(#bindings),* } => {
                            #enum_ast_path;


                            #(#stmts)*

                            #path {
                                #(#reconstruct),*
                            }
                        })
                    }
                }
            }
            Fields::Unnamed(u) => {
                let mut stmts: Vec<Stmt> = Vec::new();
                let mut bindings = Vec::<TokenStream>::new();
                let mut reconstruct = match self.kind {
                    TraitKind::Visit | TraitKind::VisitMut => None,
                    TraitKind::Fold => Some(Vec::<TokenStream>::new()),
                };

                for (idx, field) in u.unnamed.iter().enumerate() {
                    let field_name = Ident::new(&format!("_field_{}", idx), Span::call_site());
                    let ty = &field.ty;
                    let binding_idx = Lit::Int(LitInt::new(&idx.to_string(), Span::call_site()));
                    bindings.push(parse_quote!(#binding_idx: #field_name));

                    if let Some(reconstructor) = &mut reconstruct {
                        if !self.should_skip(ty) {
                            stmts.push(parse_quote!(
                                let #field_name = <#ty as #with_visitor_trait_name<V>>::#visit_with_name(#field_name, visitor #ast_path_arg);
                            ));
                        }

                        reconstructor.push(parse_quote!(#binding_idx: #field_name));
                    } else if !self.should_skip(ty) {
                        stmts.push(parse_quote!(
                            <#ty as #with_visitor_trait_name<V>>::#visit_with_name(#field_name, visitor #ast_path_arg);
                        ));
                    }
                }

                match self.kind {
                    TraitKind::Visit | TraitKind::VisitMut => {
                        parse_quote!(#path { #(#bindings),* }=> {
                            #enum_ast_path;

                            #(#stmts)*
                        })
                    }
                    TraitKind::Fold => {
                        let reconstruct = reconstruct.unwrap();

                        parse_quote!(#path { #(#bindings),* } => {
                            #enum_ast_path;

                            #(#stmts)*

                            #path{#(#reconstruct),*}
                        })
                    }
                }
            }
            Fields::Unit => match self.kind {
                TraitKind::Visit | TraitKind::VisitMut => {
                    parse_quote!(#path => {})
                }
                TraitKind::Fold => parse_quote!(#path => #path,),
            },
        }
    }

    fn implement_visit_with_for_non_node_types(&self, non_leaf_types: &[FieldType]) -> Vec<Item> {
        let visitor_trait_name = self.trait_name(false);
        let visit_with_trait_name = self.trait_name(true);
        let attrs = self.base_trait_attrs();
        let lifetime = self.method_lifetime();
        let ast_path_arg = self.arg_extra_token();
        let ast_path_param = self.param_extra_token();
        let return_type = self.return_type_token(quote!(Self));

        let receiver = self.parameter_type_token(quote!(self));

        let mut items: Vec<Item> = Vec::new();

        for node_type in non_leaf_types {
            let visit_with_name = Ident::new(
                &format!(
                    "{}_with{}",
                    self.kind.method_prefix(),
                    self.variant.method_suffix(false)
                ),
                Span::call_site(),
            );
            let visit_with_children_name = Ident::new(
                &format!(
                    "{}_children_with{}",
                    self.kind.method_prefix(),
                    self.variant.method_suffix(false)
                ),
                Span::call_site(),
            );

            let visit_method_name = Ident::new(
                &format!(
                    "{}_{}{}",
                    self.kind.method_prefix(),
                    node_type.method_name(),
                    self.variant.method_suffix(true)
                ),
                Span::call_site(),
            );

            let visit_with_doc = doc(&format!(
                "Calls [{visitor_trait_name}`::{}`] with `self`. (Extra impl)",
                visit_method_name
            ));

            let default_body: Expr = match node_type {
                FieldType::Normal(..) => match self.kind {
                    TraitKind::Visit => {
                        parse_quote!({})
                    }
                    TraitKind::VisitMut => {
                        parse_quote!({})
                    }
                    TraitKind::Fold => {
                        parse_quote!(self)
                    }
                },

                FieldType::Generic(name, inner) => match &**name {
                    "Vec" => {
                        let inner = inner.as_ref();
                        let inner_ty = quote!(#inner);

                        match (self.kind, self.variant) {
                            (TraitKind::Visit, Variant::Normal) => {
                                parse_quote!(self.iter().for_each(|item| {
                                    <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor #ast_path_arg)
                                }))
                            }
                            (TraitKind::Visit, Variant::AstPath) => {
                                parse_quote!(self.iter().enumerate().for_each(|(__idx, item)| {
                                    let mut __ast_path = __ast_path.with_index_guard(__idx);
                                    <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor, &mut *__ast_path)
                                }))
                            }
                            (TraitKind::VisitMut, Variant::Normal) => {
                                parse_quote!(
                                    self.iter_mut().for_each(|item| {
                                        <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor #ast_path_arg)
                                    })
                                )
                            }
                            (TraitKind::VisitMut, Variant::AstPath) => {
                                parse_quote!(
                                    self.iter_mut().enumerate().for_each(|(__idx, item)| {
                                        let mut __ast_path = __ast_path.with_index_guard(__idx);
                                        <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor, &mut *__ast_path)
                                    })
                                )
                            }
                            (TraitKind::Fold, Variant::Normal) => {
                                parse_quote!(
                                    swc_visit::util::move_map::MoveMap::move_map(self, |item| {
                                        <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor #ast_path_arg)
                                    })
                                )
                            }
                            (TraitKind::Fold, Variant::AstPath) => {
                                parse_quote!(
                                    self.into_iter().enumerate().map(|(__idx, item)| {
                                        let mut __ast_path = __ast_path.with_index_guard(__idx);
                                        <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(item, visitor, &mut *__ast_path)
                                    }).collect()
                                )
                            }
                        }
                    }
                    "Option" => {
                        let inner = inner.as_ref();
                        let inner_ty = quote!(#inner);

                        match self.kind {
                            TraitKind::Visit => {
                                parse_quote!(
                                    match self {
                                        Some(inner) => {
                                            <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(inner, visitor #ast_path_arg)
                                        }
                                        None => {}
                                    }
                                )
                            }
                            TraitKind::VisitMut => {
                                parse_quote!(
                                    match self {
                                        Some(inner) => {
                                            <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(inner, visitor #ast_path_arg)
                                        }
                                        None => {}
                                    }
                                )
                            }
                            TraitKind::Fold => {
                                parse_quote!(
                                    self.map(|inner| {
                                        <#inner_ty as #visit_with_trait_name<V>>::#visit_with_name(inner, visitor #ast_path_arg)
                                    })
                                )
                            }
                        }
                    }
                    "Box" => continue,
                    _ => unreachable!("unexpected generic type: {}", name),
                },
            };

            let target_type = self.node_type_for_visitor_method(node_type);

            items.push(parse_quote!(
                #(#attrs)*
                impl<V: ?Sized + #visitor_trait_name> #visit_with_trait_name<V> for #target_type {
                    #visit_with_doc
                    #[inline]
                    fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
                        <V as #visitor_trait_name>::#visit_method_name(visitor, self #ast_path_arg)
                    }

                    #[inline]
                    fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
                        #default_body
                    }
                }
            ));
        }

        items
    }

    fn implement_visit_with_for_generic_types(&self) -> Vec<Item> {
        let visit_trait_name = self.trait_name(false);
        let visit_with_trait_name = self.trait_name(true);
        let lifetime = self.method_lifetime();
        let ast_path_arg = self.arg_extra_token();
        let ast_path_param = self.param_extra_token();
        let return_type = self.return_type_token(quote!(Self));
        let attrs = self.base_trait_attrs();
        let receiver = self.parameter_type_token(quote!(self));

        let visit_with_name = Ident::new(
            &format!(
                "{}_with{}",
                self.kind.method_prefix(),
                self.variant.method_suffix(false)
            ),
            Span::call_site(),
        );
        let visit_with_children_name = Ident::new(
            &format!(
                "{}_children_with{}",
                self.kind.method_prefix(),
                self.variant.method_suffix(false)
            ),
            Span::call_site(),
        );

        let mut items = Vec::<Item>::new();

        {
            // Box<T> => T
            match self.kind {
                TraitKind::Fold => {
                    items.push(parse_quote!(
                        #(#attrs)*
                        impl<V, T> #visit_with_trait_name<V> for std::boxed::Box<T>
                            where V: ?Sized + #visit_trait_name,
                                T: #visit_with_trait_name<V> {
                            #[inline]
                            fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
                                swc_visit::util::map::Map::map(self, |inner| {
                                    <T as #visit_with_trait_name<V>>::#visit_with_name(inner, visitor #ast_path_arg)
                                })
                            }
                            #[inline]
                            fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
                                swc_visit::util::map::Map::map(self, |inner| {
                                    <T as #visit_with_trait_name<V>>::#visit_with_children_name(inner, visitor #ast_path_arg)
                                })
                            }
                        }
                    ));
                }

                _ => {
                    let deref_expr = match self.kind {
                        TraitKind::Visit => {
                            quote!(&**self)
                        }
                        TraitKind::VisitMut => {
                            quote!(&mut **self)
                        }
                        TraitKind::Fold => {
                            unreachable!()
                        }
                    };

                    let restore_expr = match self.kind {
                        TraitKind::Visit => {
                            quote!()
                        }
                        TraitKind::VisitMut => {
                            quote!()
                        }
                        TraitKind::Fold => {
                            unreachable!()
                        }
                    };

                    items.push(parse_quote!(
                        #(#attrs)*
                        impl<V, T> #visit_with_trait_name<V> for std::boxed::Box<T>
                            where V: ?Sized + #visit_trait_name,
                                T: #visit_with_trait_name<V> {
                            #[inline]
                            fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
                                let v = <T as #visit_with_trait_name<V>>::#visit_with_name(#deref_expr, visitor #ast_path_arg);
                                #restore_expr
                                v
                            }
                            #[inline]
                            fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
                                let v = <T as #visit_with_trait_name<V>>::#visit_with_children_name(#deref_expr, visitor #ast_path_arg);
                                #restore_expr
                                v
                            }
                        }
                    ));
                }
            }
        }

        if self.kind == TraitKind::Visit {
            // Vec<T> => [T]
            items.push(parse_quote!(
                #(#attrs)*
                impl<V, T> #visit_with_trait_name<V> for std::vec::Vec<T>
                    where V: ?Sized + #visit_trait_name,
                        [T]: #visit_with_trait_name<V> {
                    #[inline]
                    fn #visit_with_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
                        let v = <[T] as #visit_with_trait_name<V>>::#visit_with_name(self, visitor #ast_path_arg);
                        v
                    }

                    #[inline]
                    fn #visit_with_children_name #lifetime (#receiver, visitor: &mut V #ast_path_param) #return_type {
                        let v = <[T] as #visit_with_trait_name<V>>::#visit_with_children_name(self, visitor #ast_path_arg);
                        v
                    }
                }
            ));
        }

        items
    }

    fn node_type_for_visitor_method(&self, node_type: &FieldType) -> TokenStream {
        match self.kind {
            TraitKind::Visit => match node_type {
                FieldType::Generic(name, inner) if name == "Vec" => {
                    let inner_ty = quote!(#inner);

                    quote!([#inner_ty])
                }

                _ => quote!(#node_type),
            },
            _ => quote!(#node_type),
        }
    }
}

fn doc(s: &str) -> Attribute {
    parse_quote!(#[doc = #s])
}

fn field_variant(type_name: &Ident, field: &Field) -> Option<(TokenStream, Option<Arm>)> {
    if let Some(field_name) = &field.ident {
        let variant_name = Ident::new(&field_name.to_string().to_pascal_case(), Span::call_site());
        let variant_doc = doc(&format!("Represents [`{type_name}::{field_name}`]"));

        if extract_vec(extract_generic("Option", &field.ty).unwrap_or(&field.ty)).is_some() {
            let v = quote!(
                #variant_doc
                #variant_name(usize)
            );
            let arg = parse_quote!(
                Self::#variant_name(idx) => {
                    assert_initial_index(*idx, index);

                    *idx = index;
                },
            );
            return Some((v, Some(arg)));
        }

        return Some((
            quote!(
                #variant_doc
                #variant_name
            ),
            None,
        ));
    }

    None
}

fn extract_vec(ty: &Type) -> Option<&Type> {
    extract_generic("Vec", ty)
}

fn extract_generic<'a>(name: &str, ty: &'a Type) -> Option<&'a Type> {
    if let Type::Path(p) = ty {
        let last = p.path.segments.last().unwrap();

        if !last.arguments.is_empty() && last.ident == name {
            match &last.arguments {
                PathArguments::AngleBracketed(tps) => {
                    let arg = tps.args.first().unwrap();

                    match arg {
                        GenericArgument::Type(arg) => return Some(arg),
                        _ => unimplemented!("generic parameter other than type"),
                    }
                }
                _ => unimplemented!("Box() -> T or Box without a type parameter"),
            }
        }
    }

    if let Type::Reference(r) = ty {
        return extract_generic(name, &r.elem);
    }

    None
}

fn to_iter(e: TokenStream, ty: &Type, node_names: &[Ident]) -> Option<Expr> {
    if let Some(ty) = extract_vec(ty) {
        let inner_expr = to_iter(quote!(item), ty, node_names)?;
        return Some(parse_quote!(#e.iter().flat_map(|item| #inner_expr)));
    }

    if let Some(ty) = extract_generic("Option", ty) {
        let inner_expr = to_iter(quote!(item), ty, node_names)?;
        return Some(parse_quote!(#e.iter().flat_map(|item| #inner_expr)));
    }

    if let Some(ty) = extract_generic("Box", ty) {
        let inner_expr = to_iter(quote!(item), ty, node_names)?;
        return Some(parse_quote!({
            let item = &*#e;
            #inner_expr
        }));
    }

    if let Type::Path(p) = ty {
        let ty = &p.path.segments.last().unwrap().ident;

        if node_names.contains(ty) {
            return Some(parse_quote!(::std::iter::once(NodeRef::#ty(&#e))));
        }

        None
    } else {
        todo!("to_iter for {:?}", ty);
    }
}

fn define_fields(crate_name: &Ident, node_types: &[&Item]) -> Vec<Item> {
    let mut items = Vec::<Item>::new();
    let mut kind_enum_members = Vec::new();
    let mut parent_enum_members = Vec::new();
    let mut node_ref_enum_members = Vec::new();

    let mut kind_set_index_arms = Vec::<Arm>::new();
    let mut node_ref_set_index_arms = Vec::<Arm>::new();
    let mut node_ref_kind_arms = Vec::<Arm>::new();
    let mut node_ref_iter_next_arms = Vec::<Arm>::new();

    let node_names = node_types
        .iter()
        .filter_map(|ty| match ty {
            Item::Enum(data) => Some(data.ident.clone()),
            Item::Struct(data) => Some(data.ident.clone()),
            _ => None,
        })
        .collect::<Vec<_>>();

    let is_node_ref_raw = |ty: &Type| match ty {
        Type::Path(p) => node_names.contains(&p.path.segments.last().unwrap().ident),
        _ => false,
    };

    let is_node_ref = |ty: &Type| {
        if let Some(ty) = extract_generic("Box", ty) {
            return is_node_ref_raw(ty);
        }

        is_node_ref_raw(ty)
    };

    {
        let mut defs = Vec::<Item>::new();

        defs.push(parse_quote!(
            use #crate_name::*;
        ));

        defs.push(parse_quote!(
            #[inline(always)]
            fn assert_initial_index(idx: usize, index: usize) {
                #[cfg(debug_assertions)]
                if !(idx == usize::MAX || index == usize::MAX) {
                    {
                        panic!("Should be usize::MAX");
                    }
                }
            }
        ));

        for ty in node_types {
            let type_name = match ty {
                Item::Enum(data) => data.ident.clone(),
                Item::Struct(data) => data.ident.clone(),
                _ => continue,
            };

            let fields_enum_name = Ident::new(&format!("{type_name}Field"), Span::call_site());

            let mut variants = Vec::new();

            kind_set_index_arms.push(parse_quote!(
                Self::#type_name(v) => v.set_index(index),
            ));
            node_ref_kind_arms.push(parse_quote!(
                Self::#type_name(_, __field_kind) => AstParentKind::#type_name(*__field_kind),
            ));
            node_ref_set_index_arms.push(parse_quote!(
                Self::#type_name(_, __field_kind) => __field_kind.set_index(index),
            ));

            match ty {
                Item::Enum(data) => {
                    for variant in &data.variants {
                        let orig_ident = &variant.ident;
                        let variant_name = Ident::new(
                            &variant.ident.to_string().to_pascal_case(),
                            Span::call_site(),
                        );

                        let variant_doc = doc(&format!("Represents [`{type_name}::{orig_ident}`]"));
                        variants.push(quote!(
                            #variant_doc
                            #variant_name
                        ));
                    }

                    kind_enum_members.push(quote!(
                        #type_name(#fields_enum_name)
                    ));

                    parent_enum_members.push(quote!(
                        #type_name(&'ast #type_name, #fields_enum_name)
                    ));

                    node_ref_enum_members.push(quote!(
                        #type_name(&'ast #type_name)
                    ));

                    items.push(parse_quote!(
                        impl<'ast> From<&'ast #type_name> for NodeRef<'ast> {
                            fn from(node: &'ast #type_name) -> Self {
                                NodeRef::#type_name(node)
                            }
                        }
                    ));

                    {
                        let mut arms = Vec::<Arm>::new();

                        for variant in &data.variants {
                            let variant_name = &variant.ident;

                            // TODO: Support all kinds of fields
                            if variant.fields.len() != 1 {
                                continue;
                            }

                            for f in variant.fields.iter().filter(|f| is_node_ref(&f.ty)) {
                                let mut ty = &f.ty;
                                if let Some(inner) = extract_generic("Box", ty) {
                                    ty = inner;
                                }

                                arms.push(parse_quote!(
                                    #type_name::#variant_name(v0) => {
                                        Box::new(::std::iter::once(NodeRef::#ty(v0)))
                                    },
                                ));
                            }
                        }

                        node_ref_iter_next_arms.push(parse_quote!(
                            NodeRef::#type_name(node) => {
                                match node {
                                    #(#arms)*

                                    _ => Box::new(::std::iter::empty::<NodeRef<'ast>>())
                                }
                            }
                        ));
                    }

                    defs.push(parse_quote!(
                        impl #fields_enum_name {
                            #[inline(always)]
                            pub(crate) fn set_index(&mut self, _: usize) {
                                swc_visit::wrong_ast_path();
                            }
                        }
                    ));
                }
                Item::Struct(data) => {
                    let mut set_index_arms = Vec::<Arm>::new();

                    for field in &data.fields {
                        let opt = field_variant(&type_name, field);
                        let opt = match opt {
                            Some(v) => v,
                            None => continue,
                        };
                        variants.push(opt.0);

                        set_index_arms.extend(opt.1);
                    }

                    kind_enum_members.push(quote!(
                        #type_name(#fields_enum_name)
                    ));

                    parent_enum_members.push(quote!(
                        #type_name(&'ast #type_name, #fields_enum_name)
                    ));

                    node_ref_enum_members.push(quote!(
                        #type_name(&'ast #type_name)
                    ));

                    items.push(parse_quote!(
                        impl<'ast> From<&'ast #type_name> for NodeRef<'ast> {
                            fn from(node: &'ast #type_name) -> Self {
                                NodeRef::#type_name(node)
                            }
                        }
                    ));

                    {
                        let mut iter: Expr = parse_quote!(::std::iter::empty::<NodeRef<'ast>>());

                        match &data.fields {
                            Fields::Named(fields) => {
                                for f in fields.named.iter() {
                                    let ident = &f.ident;
                                    let iter_expr =
                                        to_iter(quote!(node.#ident), &f.ty, &node_names);
                                    if let Some(iter_expr) = iter_expr {
                                        iter = parse_quote!(#iter.chain(#iter_expr));
                                    }
                                }
                            }

                            Fields::Unnamed(_fields) => {
                                // TODO: Support unnamed fields
                            }
                            Fields::Unit => {}
                        }

                        node_ref_iter_next_arms.push(parse_quote!(
                            NodeRef::#type_name(node) => {
                                let iterator = #iter;
                                Box::new(iterator)
                            }
                        ));
                    }

                    defs.push(parse_quote!(
                        impl #fields_enum_name {
                            pub(crate) fn set_index(&mut self, index: usize) {
                                match self {
                                    #(#set_index_arms)*

                                    _ => {
                                        swc_visit::wrong_ast_path()
                                    }
                                }
                            }
                        }
                    ));
                }
                _ => continue,
            }

            defs.push(parse_quote!(
                #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
                #[cfg_attr(feature = "serde-impl", derive(serde::Serialize, serde::Deserialize))]
                pub enum #fields_enum_name {
                    #(#variants),*
                }
            ))
        }

        {
            defs.push(parse_quote!(
                #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
                #[cfg_attr(feature = "serde-impl", derive(serde::Serialize, serde::Deserialize))]
                pub enum AstParentKind {
                    #(#kind_enum_members),*
                }
            ));

            defs.push(parse_quote!(
                impl ::swc_visit::ParentKind for AstParentKind {
                    #[inline]
                    fn set_index(&mut self, index: usize) {
                        match self {
                            #(#kind_set_index_arms)*
                        }
                    }
                }
            ));
        }

        {
            defs.push(parse_quote!(
                #[derive(Debug, Clone, Copy)]
                pub enum AstParentNodeRef<'ast> {
                    #(#parent_enum_members),*
                }
            ));
            items.push(parse_quote!(
                #[derive(Debug, Clone, Copy)]
                pub enum NodeRef<'ast> {
                    #(#node_ref_enum_members),*
                }
            ));

            defs.push(parse_quote!(
                impl<'ast> ::swc_visit::NodeRef for AstParentNodeRef<'ast> {
                    type ParentKind = AstParentKind;

                    #[inline(always)]
                    fn kind(&self) -> AstParentKind {
                        self.kind()
                    }

                    fn set_index(&mut self, index: usize) {
                        match self {
                            #(#node_ref_set_index_arms)*
                        }
                    }
                }
            ));
            defs.push(parse_quote!(
                #[cfg(any(docsrs, feature = "path"))]
                impl<'ast> AstParentNodeRef<'ast> {
                    #[inline]
                    pub fn kind(&self) -> AstParentKind {
                        match self {
                            #(#node_ref_kind_arms)*
                        }
                    }
                }
            ));
            items.push(parse_quote!(
                impl<'ast> NodeRef<'ast> {
                    /// This is not a part of semver-stable API. It is experimental and subject to change.
                    #[allow(unreachable_patterns)]
                    pub fn experimental_raw_children<'a>(&'a self) -> Box<dyn 'a + Iterator<Item = NodeRef<'ast>>> {
                        match self {
                            #(#node_ref_iter_next_arms)*
                        }
                    }
                }
            ));

            items.push(parse_quote!(
                impl<'ast> NodeRef<'ast> {
                    /// Visit all nodes in self in preorder.
                    ///
                    /// This is not a part of semver-stable API. It is
                    /// experimental and subject to change.
                    pub fn experimental_traverse(
                        &'ast self,
                    ) -> Box<dyn 'ast + Iterator<Item = NodeRef<'ast>>> {
                        let mut queue = std::collections::VecDeque::<NodeRef<'ast>>::new();
                        queue.push_back(*self);

                        Box::new(std::iter::from_fn(move || {
                            let node: NodeRef<'ast> = queue.pop_front()?;
                            {
                                let children = node.experimental_raw_children();
                                queue.extend(children);
                            }
                            Some(node)
                        }))
                    }
                }
            ));
        }

        items.insert(
            0,
            parse_quote!(
                #[cfg(any(docsrs, feature = "path"))]
                pub mod fields {
                    #(#defs)*
                }
            ),
        );

        items.push(parse_quote!(
            #[cfg(any(docsrs, feature = "path"))]
            pub use self::fields::{AstParentKind, AstParentNodeRef};
        ));
    }

    items
}
