use proc_macro::TokenStream;
use quote::{format_ident, quote, ToTokens};
use syn::{
    parse::Nothing, parse_macro_input, FnArg, Generics, ItemTrait, PatType, Receiver, Signature,
    TraitItem, TraitItemFn,
};

#[proc_macro_attribute]
pub fn ffi_trait(input: TokenStream, item: TokenStream) -> TokenStream {
    parse_macro_input!(input as Nothing);

    let mut tokens = proc_macro2::TokenStream::from(item.clone());
    let t = parse_macro_input!(item as ItemTrait);
    let ItemTrait {
        attrs: _,
        vis,
        unsafety: _,
        auto_token: _,
        restriction: _,
        trait_token: _,
        ident: trait_ident,
        generics: _trait_generics,
        colon_token: _,
        supertraits: _,
        brace_token: _,
        items,
    } = &t;

    let (_trait_impl_generics, _trait_ty_generics, trait_where_clause) =
        _trait_generics.split_for_impl();

    let _trait_impl_generics = syn::parse2::<Generics>(_trait_impl_generics.to_token_stream())
        .unwrap()
        .params;
    let mut trait_impl_generics = _trait_impl_generics.to_token_stream();
    if !_trait_impl_generics.empty_or_trailing() {
        trait_impl_generics.extend(quote! { , });
    }

    let _trait_ty_generics = syn::parse2::<Generics>(_trait_ty_generics.to_token_stream())
        .unwrap()
        .params;
    let mut trait_ty_generics = _trait_ty_generics.to_token_stream();
    if !_trait_ty_generics.empty_or_trailing() {
        trait_ty_generics.extend(quote! { , });
    }

    let ffi_dyn_trait_name = format_ident!("FfiDyn{trait_ident}");
    let ffi_dyn_trait_vtable_name = format_ident!("FfiDyn{trait_ident}Vtable");
    let ffi_dyn_trait_self_name = format_ident!("FfiDyn{trait_ident}Self");

    let mut impl_trait_tokens = proc_macro2::TokenStream::new();
    let mut vtable_functions = proc_macro2::TokenStream::new();
    let mut thunk_members = proc_macro2::TokenStream::new();
    for item in items {
        let TraitItem::Fn(TraitItemFn {
            attrs: _,
            sig:
                Signature {
                    constness: _,
                    asyncness: None,
                    unsafety: _,
                    abi: _,
                    fn_token: _,
                    ident,
                    generics,
                    paren_token: _,
                    inputs,
                    variadic: None,
                    output,
                },
            default: _,
            semi_token: _,
        }) = item
        else {
            todo!("only trait methods are allowed")
        };

        let mut has_self = false;
        let mut impl_trait_params = proc_macro2::TokenStream::new();
        let mut vtable_params = proc_macro2::TokenStream::new();
        let mut thunk_params = proc_macro2::TokenStream::new();
        let mut impl_trait_args = proc_macro2::TokenStream::new();
        let mut thunk_args = proc_macro2::TokenStream::new();
        for (i, input) in inputs.iter().enumerate() {
            match input {
                FnArg::Receiver(receiver) => match receiver {
                    Receiver {
                        attrs: _,
                        reference: Some((_, lifetime)),
                        mutability,
                        self_token: _,
                        colon_token: _,
                        ty: _,
                    } => {
                        has_self = true;
                        impl_trait_params.extend(quote! { &#lifetime #mutability self, });
                        vtable_params.extend(quote! { this: #ffi_dyn_trait_self_name<#lifetime>, });
                        thunk_params.extend(quote! { this: &#lifetime #mutability _Self, });
                        impl_trait_args.extend(quote! { self.ptr, });
                        thunk_args.extend(quote! { this, });
                    }
                    _ => todo!(),
                },
                FnArg::Typed(PatType {
                    attrs: _,
                    pat: _,
                    colon_token: _,
                    ty,
                }) => {
                    let param_name = format_ident!("_{i}");
                    impl_trait_params.extend(quote! { #param_name: #ty, });
                    vtable_params.extend(quote! { #param_name: #ty, });
                    thunk_params.extend(quote! { #param_name: #ty, });
                    impl_trait_args.extend(quote! { #param_name, });
                    thunk_args.extend(quote! { #param_name, });
                }
            }
        }

        let mut no_args_tokens = proc_macro2::TokenStream::new();
        if !has_self {
            no_args_tokens.extend(quote! { &self });
        }

        impl_trait_tokens.extend(quote! {
            #vis fn #ident #generics(#no_args_tokens #impl_trait_params) #output {
                unsafe { ((*self.vtable).#ident)(#impl_trait_args) }
            }
        });
        let generics_without_braces = &generics.params;
        vtable_functions.extend(quote! {
            #ident: for<#generics_without_braces> unsafe extern "C" fn(#vtable_params) #output,
        });
        let mut generics_without_braces_tokens = generics_without_braces.to_token_stream();
        if !generics_without_braces.empty_or_trailing() {
            generics_without_braces_tokens.extend(quote! { , });
        }
        thunk_members.extend(quote! {
            #ident: {
                unsafe extern "C" fn thunk<#generics_without_braces_tokens #trait_impl_generics _Self: #trait_ident<#trait_ty_generics>>(#thunk_params) #output #trait_where_clause {
                    unsafe { <_Self as #trait_ident<#trait_ty_generics>>::#ident(#thunk_args) }
                }
                #[allow(clippy::useless_transmute)]
                unsafe { ::core::mem::transmute(thunk::<#trait_ty_generics _Self> as for<#generics_without_braces> unsafe extern "C" fn(#thunk_params) #output) }
            },
        });
    }

    tokens.extend(quote! {
        #[derive(::core::clone::Clone, ::core::marker::Copy)]
        #[repr(transparent)]
        struct #ffi_dyn_trait_self_name<'_self> {
            ptr: ::core::ptr::NonNull<()>,
            _data: ::core::marker::PhantomData<&'_self ()>,
        }

        #[repr(C)]
        #vis struct #ffi_dyn_trait_name<'_self, #trait_impl_generics> #trait_where_clause {
            ptr: #ffi_dyn_trait_self_name<'_self>,
            vtable: *const #ffi_dyn_trait_vtable_name<#trait_ty_generics>,
            _generics: ::core::marker::PhantomData<(#trait_ty_generics)>,
        }

        impl<'_self, #trait_impl_generics> #ffi_dyn_trait_name<'_self, #trait_ty_generics> #trait_where_clause {
            #vis fn new<_Self: #trait_ident<#trait_ty_generics>>(value: &'_self mut _Self) -> Self {
                Self {
                    ptr: #ffi_dyn_trait_self_name {
                        ptr: ::core::ptr::NonNull::from(value).cast(),
                        _data: ::core::marker::PhantomData,
                    },
                    vtable: &const { <#ffi_dyn_trait_vtable_name<#trait_ty_generics>>::new::<_Self>() },
                    _generics: ::core::marker::PhantomData,
                }
            }
        }

        impl<'_self, #trait_impl_generics> #ffi_dyn_trait_name<'_self, #trait_ty_generics> #trait_where_clause {
            #impl_trait_tokens
        }

        #[repr(C)]
        struct #ffi_dyn_trait_vtable_name<#trait_impl_generics> #trait_where_clause {
            #vtable_functions
            _generics: ::core::marker::PhantomData<(#trait_ty_generics)>,
        }

        impl<#trait_impl_generics> #ffi_dyn_trait_vtable_name<#trait_ty_generics> #trait_where_clause {
            #vis const fn new<_Self: #trait_ident<#trait_ty_generics>>() -> Self {
                Self {
                    #thunk_members
                    _generics: ::core::marker::PhantomData,
                }
            }
        }
    });

    tokens.into()
}
