// SPDX-FileCopyrightText: 2021 - 2023 Kamila Borowska <kamila@borowska.pw>
// SPDX-FileCopyrightText: 2021 Bruno Corrêa Zimmermann <brunoczim@gmail.com>
//
// SPDX-License-Identifier: MIT OR Apache-2.0

use crate::type_length;
use proc_macro2::TokenStream;
use quote::{format_ident, quote, ToTokens};
use syn::{DataEnum, Fields, FieldsNamed, FieldsUnnamed, Ident, Variant};

pub fn generate(name: Ident, data_enum: DataEnum) -> TokenStream {
    let mut generator = EnumGenerator::empty();
    for variant in &data_enum.variants {
        generator.handle_variant(variant);
    }
    generator.finish(&name)
}

#[derive(Debug)]
struct Length {
    units: usize,
    opaque: TokenStream,
}

impl ToTokens for Length {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let Self { units, opaque } = self;
        tokens.extend(quote! { (#units + #opaque) });
    }
}

/// Total length is the sum of each variant's length. To represent a variant, its number is added to
/// the sum of previous variant lengths.
#[derive(Debug)]
struct EnumGenerator {
    length: Length,
    from_usize_arms: TokenStream,
    into_usize_arms: TokenStream,
}

impl EnumGenerator {
    fn empty() -> Self {
        Self {
            length: Length {
                units: 0,
                opaque: quote! { 0 },
            },
            from_usize_arms: quote! {},
            into_usize_arms: quote! {},
        }
    }

    fn finish(&self, name: &Ident) -> TokenStream {
        let Self {
            length,
            from_usize_arms,
            into_usize_arms,
        } = self;

        quote! {
            #[automatically_derived]
            impl ::enum_map::Enum for #name {
                const LENGTH: ::enum_map::usize = #length;

                #[inline]
                fn from_usize(value: ::enum_map::usize) -> Self {
                    #from_usize_arms {
                        ::enum_map::out_of_bounds()
                    }
                }

                #[inline]
                fn into_usize(self) -> ::enum_map::usize {
                    match self {
                        #into_usize_arms
                    }
                }
            }

            #[automatically_derived]
            impl<V> ::enum_map::EnumArray<V> for #name {
                type Array = [V; #length];
            }
        }
    }

    fn handle_variant(&mut self, variant: &Variant) {
        match &variant.fields {
            Fields::Unit => self.handle_unit_variant(&variant.ident),
            Fields::Unnamed(fields) => self.handle_unnamed_variant(&variant.ident, fields),
            Fields::Named(fields) => self.handle_named_variant(&variant.ident, fields),
        }
    }

    /// Becomes simply `1` in counting, since this is the size of the unit.
    fn handle_unit_variant(&mut self, variant: &Ident) {
        let Self {
            length,
            from_usize_arms,
            into_usize_arms,
        } = self;
        *into_usize_arms = quote! { #into_usize_arms Self::#variant => #length, };
        *from_usize_arms = quote! {
            #from_usize_arms if value == #length {
                Self::#variant
            } else
        };
        self.length.units += 1;
    }

    /// Its size is the product of the sizes of its members. To represent this variant, one can
    /// think of this as representing a little-endian number. First member is simply added, but
    /// next members are multiplied before being added.
    fn handle_unnamed_variant(&mut self, variant: &Ident, fields: &FieldsUnnamed) {
        let Self {
            length,
            from_usize_arms,
            into_usize_arms,
        } = self;
        let mut expr_into = quote! { #length };
        let mut fields_length = quote! { 1usize };
        let mut params_from = quote! {};
        for (i, field) in fields.unnamed.iter().enumerate() {
            let ident = format_ident!("p{}", i);
            let ty = &field.ty;
            let field_length = type_length(ty);

            expr_into = quote! {
                (#expr_into + #fields_length * ::enum_map::Enum::into_usize(#ident))
            };

            params_from = quote! {
                #params_from <#ty as ::enum_map::Enum>::from_usize(
                    (value - #length) / #fields_length % #field_length
                ),
            };

            fields_length = quote! { (#fields_length * #field_length) };
        }

        *length = Length {
            units: 0,
            opaque: quote! { (#length + #fields_length) },
        };

        let from_arms = &from_usize_arms;
        *from_usize_arms = quote! {
            #from_arms if value < #length {
                Self::#variant(#params_from)
            } else
        };

        let mut params_into = quote! {};
        for i in 0..fields.unnamed.len() {
            let ident = format_ident!("p{}", i);
            params_into = quote! { #params_into #ident, };
        }

        *into_usize_arms = quote! {
            #into_usize_arms Self::#variant(#params_into) => #expr_into,
        };
    }

    /// Its size is the product of the sizes of its members. To represent this variant, one can
    /// think of this as representing a little-endian number. First member is simply added, but
    /// next members are multiplied before being added.
    fn handle_named_variant(&mut self, variant: &Ident, fields: &FieldsNamed) {
        let Self {
            length,
            from_usize_arms,
            into_usize_arms,
        } = self;
        let mut expr_into = quote! { #length };
        let mut fields_length = quote! { 1usize };
        let mut params_from = quote! {};

        for field in fields.named.iter() {
            let ident = field.ident.as_ref().unwrap();
            let ty = &field.ty;
            let field_length = type_length(ty);

            expr_into = quote! {
                (#expr_into + #fields_length * ::enum_map::Enum::into_usize(#ident))
            };

            params_from = quote! {
                #params_from #ident: <#ty as ::enum_map::Enum>::from_usize(
                    (value - #length) / #fields_length % #field_length
                ),
            };

            fields_length = quote! { (#fields_length * #field_length) };
        }

        *length = Length {
            units: 0,
            opaque: quote! { (#length + #fields_length) },
        };

        *from_usize_arms = quote! {
            #from_usize_arms if value < #length {
                Self::#variant { #params_from }
            } else
        };

        let mut params_into = quote! {};
        for field in fields.named.iter() {
            let ident = field.ident.as_ref().unwrap();
            params_into = quote! { #params_into #ident, };
        }

        *into_usize_arms = quote! {
            #into_usize_arms Self::#variant { #params_into } => #expr_into,
        };
    }
}