#![deny(clippy::all, clippy::pedantic)]

use proc_macro::TokenStream;
use quote::{ToTokens, format_ident, quote};
use syn::{Fields, Item, ItemEnum, ItemStruct, LitStr, Path as SynPath, parse};

macro_rules! syn_path {
    ($first_segment: ident $(::$segment: ident)*) => {
        ::syn::Path {
            leading_colon: None,
            segments: ::syn::punctuated::Punctuated::from_iter([
                syn_path_segment!($first_segment),
                $(syn_path_segment!($segment),)*
            ])
        }
    };
}

macro_rules! syn_path_segment {
    ($segment: ident) => {
        ::syn::PathSegment {
            ident: ::proc_macro2::Ident::new(
                stringify!($segment),
                ::proc_macro2::Span::call_site(),
            ),
            arguments: ::syn::PathArguments::None,
        }
    };
}

#[proc_macro_derive(Reflection)]
pub fn reflection_derive(input: TokenStream) -> TokenStream
{
    let input = parse::<Item>(input).unwrap();

    match input {
        Item::Struct(input) => generate_struct_reflection_impl(input),
        Item::Enum(input) => generate_enum_reflection_impl(input),
        _ => panic!("Invalid input"),
    }
}

fn generate_struct_reflection_impl(input: ItemStruct) -> TokenStream
{
    let engine_crate_path = find_engine_crate_path().unwrap();

    let input_ident = input.ident;

    let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();

    let fields = input.fields.into_iter().enumerate().map(|(index, field)| {
        let field_ident = field.ident.unwrap_or_else(|| format_ident!("{index}"));

        let field_type = &field.ty;

        let field_name = LitStr::new(&field_ident.to_string(), field_ident.span());

        // since std::any::type_name as const is not stable yet
        let field_type_name = field_type.to_token_stream().to_string();

        quote! {
            #engine_crate_path::reflection::StructField {
                name: #field_name,
                index: #index,
                layout: std::alloc::Layout::new::<#field_type>(),
                byte_offset: std::mem::offset_of!(Self, #field_ident),
                type_id: std::any::TypeId::of::<#field_type>(),
                type_name: #field_type_name,
                get_type: #engine_crate_path::reflection::FnWithDebug::new(|| {
                    struct SpecializationTarget<Field>(std::marker::PhantomData<Field>);

                    trait FieldHasReflection
                    {
                        fn field_type_reflection(&self)
                            -> Option<&'static #engine_crate_path::reflection::Type>;
                    }

                    trait FieldDoesNotHaveReflection
                    {
                        fn field_type_reflection(&self)
                            -> Option<&'static #engine_crate_path::reflection::Type>;
                    }

                    impl<Field> FieldDoesNotHaveReflection for &SpecializationTarget<Field>
                    {
                        fn field_type_reflection(&self)
                            -> Option<&'static #engine_crate_path::reflection::Type>
                        {
                            None
                        }
                    }

                    impl<Field> FieldHasReflection for SpecializationTarget<Field>
                    where
                        Field: #engine_crate_path::reflection::Reflection
                    {
                        fn field_type_reflection(&self)
                            -> Option<&'static #engine_crate_path::reflection::Type>
                        {
                            Some(Field::type_reflection())
                        }
                    }

                    (&SpecializationTarget::<#field_type>(std::marker::PhantomData))
                        .field_type_reflection()
                })
            }
        }
    });

    quote! {
        unsafe impl #impl_generics #engine_crate_path::reflection::Reflection for
            #input_ident #type_generics #where_clause
        {
            const TYPE_REFLECTION: &#engine_crate_path::reflection::Type =
                &const {
                    #engine_crate_path::reflection::Type::Struct(
                        #engine_crate_path::reflection::Struct {
                            fields: &[
                                #(#fields),*
                            ]
                        }
                    )
                };
        }
    }
    .into()
}

fn generate_enum_reflection_impl(input: ItemEnum) -> TokenStream
{
    let engine_crate_path = find_engine_crate_path().unwrap();

    let input_ident = input.ident;

    let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();

    let variants = input.variants.iter().map(|variant| {
        let variant_name = LitStr::new(&variant.ident.to_string(), variant.ident.span());

        quote! {
            #engine_crate_path::reflection::EnumVariant {
                name: #variant_name,
            }
        }
    });

    let variant_lookup_match_arms =
        input.variants.iter().enumerate().map(|(index, variant)| {
            let variant_ident = &variant.ident;

            let pattern = match variant.fields {
                Fields::Unit => quote! { Self::#variant_ident },
                Fields::Named(_) => quote! { Self::#variant_ident { .. } },
                Fields::Unnamed(_) => quote! { Self::#variant_ident(..) },
            };

            quote! {
                #pattern => &enum_reflection.variants[#index]
            }
        });

    quote! {
        unsafe impl #impl_generics #engine_crate_path::reflection::Reflection for
            #input_ident #type_generics #where_clause
        {
            const TYPE_REFLECTION: &#engine_crate_path::reflection::Type =
                &const {
                    #engine_crate_path::reflection::Type::Enum(
                        #engine_crate_path::reflection::Enum {
                            variants: &[
                                #(#variants),*
                            ]
                        }
                    )
                };
        }

        unsafe impl #impl_generics #engine_crate_path::reflection::EnumReflectionExt for
            #input_ident #type_generics #where_clause
        {
            fn get_variant_reflection(&self)
                -> &'static #engine_crate_path::reflection::EnumVariant
            {
                let enum_reflection = unsafe {
                    <Self as #engine_crate_path::reflection::Reflection>::TYPE_REFLECTION
                        .as_enum()
                        .unwrap_unchecked()
                };

                match self {
                    #(#variant_lookup_match_arms),*
                }
            }
        }
    }
    .into()
}

fn find_engine_crate_path() -> Option<SynPath>
{
    let cargo_crate_name = std::env::var("CARGO_CRATE_NAME").ok()?;
    let cargo_pkg_name = std::env::var("CARGO_PKG_NAME").ok()?;

    if cargo_pkg_name == "engine" && cargo_crate_name != "engine" {
        // Macro is used by a crate example/test/benchmark
        return Some(syn_path!(engine));
    }

    if cargo_crate_name == "engine" {
        return Some(syn_path!(crate));
    }

    Some(syn_path!(engine))
}