#![doc(test(attr(
    deny(warnings),
    allow(unused),
    deny(unused_unsafe),
)))]
#![cfg_attr(feature = "nightly",
    feature(external_doc),
)]
#![cfg_attr(feature = "nightly",
    doc(include = "../README.md")
)]

#[cfg(not(all(test, feature = "unit-tests")))]
extern crate proc_macro;
#[cfg(not(all(test, feature = "unit-tests")))]
use ::proc_macro::TokenStream;
#[cfg(not(all(test, feature = "unit-tests")))]
use ::syn::parse;
#[cfg(all(test, feature = "unit-tests"))]
use ::proc_macro2::TokenStream;
#[cfg(all(test, feature = "unit-tests"))]
use ::syn::parse2 as parse;

macro_rules! parse_macro_input {
    (
        $tokenstream:ident as $ty:ty
    ) => (
        match parse::<$ty>($tokenstream) {
            Ok(data) => data,
            Err(err) => {
                return TokenStream::from(err.to_compile_error());
            }
        }
    );

    (
        $tokenstream:ident
    ) => (
        parse_macro_input!($tokenstream as _)
    );
}

use ::proc_macro2::{
    Span,
    TokenStream as TokenStream2,
};
use ::quote::{
    quote,
    quote_spanned,
};
use ::syn::{*,
    parse::{
        Parse,
        ParseStream,
    },
    spanned::{
        Spanned,
    },
};
use ::std::{*,
    iter::FromIterator,
    ops::Not,
    result::Result::{self, *},
};

use self::utils::{
    Either,
    OptionInsert,
};

#[macro_use]
mod utils;

struct FunctionLike<'a> {
    attrs: &'a Vec<Attribute>,
    sig: &'a mut Signature,
    block: &'a mut Block,
}

fn edit_method_in_place (
    fn_item: FunctionLike<'_>,
    self_ty: Option<(
        &'_ Generics,
        &'_ Box<Type>,
        Option<(
            Option<&'_ TypeParam>,
            &'_ [TypeParamBound],
        )>,
    )>,
) -> Result<(), TokenStream>
{
    if self_ty.is_none() {
        if utils::is_method(&fn_item) {
            return Err(error!(concat!(
                "`#[require_unsafe_in_body]` does not support directly ",
                "decorating a method; you need to decorate the whole ",
                "`impl` or `trait` block with `#[require_unsafe_in_bodies]`",
            )));
        }
    }
    if fn_item.sig.unsafety.is_none() {
        return Ok(());
    }
    let inner: ImplItemMethod = {
        let mut inner = ImplItemMethod {
            attrs: fn_item.attrs.clone(),
            sig: fn_item.sig.clone(),
            block: fn_item.block.clone(),

            vis: Visibility::Inherited,
            defaultness: None,
        };
        inner.sig.ident = parse_quote!(
            __require_unsafe__inner
        );
        // strip unsafe <= This is the point of the whole macro
        inner.sig.unsafety = None;
        // strip extern
        inner.sig.abi = None;
        // strip (some) special attributes
        inner
            .attrs
            .retain(|&Attribute { ref path, ..}| bool::not(
                path.is_ident("inline") ||
                path.is_ident("cold") ||
                path.is_ident("no_mangle")
            ))
        ;
        inner.attrs.push(parse_quote! {
            #[inline(always)]
        });
        inner
    };
    // Append an `.await` to the forwarded call if the function is `async`
    let _await =
        inner
            .sig
            .asyncness
            .map(|_| quote!( .await ))
            .unwrap_or_default()
    ;
    let mut with_receiver = false;
    let args: Vec<Expr> =
        (0 ..)
            .map(|i| {
                let s: &str = &*format!("__require_unsafe__arg_{}", i);
                Ident::new(s, Span::call_site())
            })
            .zip(fn_item.sig.inputs.iter_mut())
            // replace input parameter patterns by the numbered `arg_name`s,
            // unless the parameter is `self`,
            // while outputing these numbered `arg_name`s.
            .map(|(arg_name, at_fn_arg)| match at_fn_arg {
                | &mut FnArg::Receiver(_) => {
                    with_receiver = true;
                    parse_quote!(self)
                },
                | &mut FnArg::Typed(PatType {
                    ref mut pat,
                    ..
                }) => {
                    if let Pat::Ident(ref pat_ident) = **pat {
                        if pat_ident.ident == "self" {
                            with_receiver = true;
                            return parse_quote!(self);
                        }
                    }
                    **pat = parse_quote!(#arg_name); // Pat
                    parse_quote!(#arg_name) // Expr
                },
            })
            .collect()
    ;
    *fn_item.block = match self_ty {
        // Classic function
        | None => {
            let ty_generics: TypeGenerics =
                fn_item
                    .sig
                    .generics
                    .split_for_impl()
                    .1
            ;
            let ty_generics: Turbofish = ty_generics.as_turbofish();
            parse_quote!({
                #inner

                __require_unsafe__inner #ty_generics (
                    #(#args),*
                ) #_await
            })
        },

        // # trait / impl function
        //
        // may use associated items that are not in scope in the helper trait,
        // hence the following strategy is used:
        //
        //  1. the receiver_argument (self), if any, is extracted,
        //
        //  2. then all the other arguments types are abstracted within helper
        //     associated types
        //
        //  3. and so is the return type.
        /*
        trait #impl_generics Helper #supertraits_bounds
        #where_clauses
        {
            type Arg0; // if no self

            #(
                type ArgN;
            )*

            type Ret;

            fn inner (#(#self,)? #(_: Self::ArgN),*) -> Self::Ret;
        }
        impl #impl_generics Helper
            for Self
        #where_clauses
        {
            #(
                type #ArgN = #args[N];
            )*

            type Ret = #Ret;

            #inner
        }
        */
        | Some((
            outer_generics,
            self_ty,
            mb_impl_trait_param,
        )) => {
            let mut inner = inner; // make mutable
            let mut outer_generics = outer_generics.clone();
            // remove the genericity of inner, by moving it into the outer
            // genericity
            utils::merge_generics(
                &mut outer_generics,
                &mem::replace(
                    &mut inner.sig.generics,
                    Generics::default(),
                ),
            )?;
            let inner = inner; // done mutating inner
            let (
                trait_generics,
                ty_generics,
                where_clause,
            ) =
                outer_generics
                    .split_for_impl()
            ;
            // Iterator of enumerated inner input type parameters
            //  \-> yields (i: usize, arg: FnArg) pairs
            let mut inputs =
                inner
                    .sig
                    .inputs
                    .iter()
                    .cloned()
                    .enumerate()
                    .peekable()
            ;
            // #(#self,)? #(_: Self::ArgN),*
            let mut helper_args =
                Vec::<FnArg>::with_capacity(0)
            ;
            match inputs.peek() {
                | Some(&(_, ref arg))
                    if utils::is_receiver(arg)
                => {
                    let (_, arg) = inputs.next().unwrap();
                    helper_args.push(arg);
                },
                | _ => {},
            }
            let (args_idents, args_tys): (Vec<Ident>, Vec<Type>) =
                inputs
                    .map(|(i, arg)| match arg {
                        | FnArg::Typed(arg) => {
                            let ident = Ident::new(
                                &format!("__require_unsafe__Arg{}", i),
                                Span::call_site(),
                            );
                            helper_args.push(parse_quote! {
                                _ : Self::#ident
                            });
                            (ident, *arg.ty)
                        },
                        | _ => unreachable!(),
                    })
                    .unzip()
            ;
            let mut helper_sig = inner.sig.clone();
            helper_sig.inputs = helper_args.into_iter().collect();
            #[allow(bad_style)]
            let Ret: Type =
                match   mem::replace(
                            &mut helper_sig.output,
                            parse_quote! {
                                -> Self::__require_unsafe__Ret
                            },
                        )
                {
                    | ReturnType::Default => parse_quote! { () },
                    | ReturnType::Type(_, ty) => *ty,
                }
            ;
            let mut supertraits_bounds = TokenStream2::new();
            // locals for lifetime extension
            let mut generics = None::<Generics>;
            let mut impl_generics = None::<ImplGenerics>;
            let impl_generics: &ImplGenerics =
                mb_impl_trait_param
                    .and_then(|(impl_trait_param, supertraits)| {
                        if supertraits.is_empty().not() {
                            supertraits_bounds = quote! {
                                : #(#supertraits)+*
                            };
                        }
                        impl_trait_param
                            .map(|impl_trait_param| -> &ImplGenerics {
                                let generics = generics.insert(
                                    outer_generics.clone()
                                );
                                generics.params.push(
                                    impl_trait_param.clone().into()
                                );
                                impl_generics.insert(
                                    generics.split_for_impl().0
                                )
                            })
                    })
                    .unwrap_or(&trait_generics)
            ;
            parse_quote!({
                trait __require_unsafe__trait #trait_generics
                #supertraits_bounds
                #where_clause
                {
                    #(
                        type #args_idents;
                    )*
                    type __require_unsafe__Ret;
                    #helper_sig;
                }
                impl #impl_generics __require_unsafe__trait #ty_generics
                    for #self_ty
                #where_clause
                {
                    #(
                        type #args_idents = #args_tys;
                    )*
                    type __require_unsafe__Ret = #Ret;
                    #inner
                }

                <Self as __require_unsafe__trait #ty_generics>
                    ::__require_unsafe__inner(
                        #(#args),*
                    ) #_await
            })
        },
    };
    Ok(())
}

#[cfg_attr(feature = "nightly",
    doc(include = "docs/require_unsafe_in_body.md")
)]
#[proc_macro_attribute] pub
fn require_unsafe_in_body (
    attrs: TokenStream,
    input: TokenStream,
) -> TokenStream
{
    let attrs = TokenStream2::from(attrs);
    if attrs.is_empty().not() {
        return error!(attrs.span()=>
            "Unexpected parameter(s)"
        );
    }
    let mut item_fn: ImplItemMethod = match parse(input.clone()) {
        | Ok(item_fn) => item_fn,
        | Err(error) => {
            if  parse::<ItemImpl>(input.clone()).is_ok() ||
                parse::<ItemTrait>(input).is_ok()
            {
                return error!(
                    concat!(
                        "To decorate an `impl` or `trait` block, ",
                        "you need to use `#[require_unsafe_in_bodies]`",
                    ),
                );
            } else {
                return error.to_compile_error().into();
            }
        },
    };
    if item_fn.sig.unsafety.is_none() {
        return input;
    }
    let ImplItemMethod {
        ref attrs,
        ref mut sig,
        ref mut block,
        ..
    } = item_fn;
    let function_like = FunctionLike {
        attrs,
        sig,
        block,
    };
    unwrap!(
        edit_method_in_place(function_like, None)
    );
    TokenStream::from(quote! {
        #item_fn
    })
}

#[cfg_attr(feature = "nightly",
    doc(include = "docs/require_unsafe_in_bodies.md")
)]
#[proc_macro_attribute] pub
fn require_unsafe_in_bodies (
    attrs: TokenStream,
    input: TokenStream,
) -> TokenStream
{
    let attrs = TokenStream2::from(attrs);
    if attrs.is_empty().not() {
        return error!(attrs.span()=>
            "Unexpected parameter(s)"
        );
    }
    match parse_macro_input!(input as Either<ItemImpl, ItemTrait>) {
        | Either::Left(mut item_impl) => {
            let mut type_param_bound: Option<TypeParamBound> = None;
            let super_trait:
                Option<(
                    Option<&TypeParam>,
                    &[TypeParamBound],
                )>
            =
                item_impl
                    .trait_
                    .as_ref()
                    .map(|(_, trait_, _)| (
                        None,
                        slice::from_ref(
                            type_param_bound
                                .get_or_insert(TraitBound {
                                    paren_token: None,
                                    modifier: TraitBoundModifier::None,
                                    lifetimes: None,
                                    path: trait_.clone(),
                                }.into())
                        ),
                    ))
            ;
            let ItemImpl {
                ref generics,
                ref self_ty,
                ref mut items,
                ..
            } = item_impl;
            unwrap!(
                items
                    .iter_mut()
                    .try_for_each(|impl_item| Ok(match *impl_item {
                        | ImplItem::Method(ImplItemMethod {
                            ref attrs,
                            ref mut sig,
                            ref mut block,
                            ..
                        }) => {
                            let function_like = FunctionLike {
                                attrs,
                                sig,
                                block,
                            };
                            edit_method_in_place(
                                function_like,
                                Some((generics, self_ty, super_trait)),
                            )?;
                        },
                        | _ => {},
                    }))
            );
            TokenStream::from(quote! {
                #item_impl
            })
        },

        | Either::Right(mut item_trait) => {
            let ItemTrait {
                ref generics,
                ref mut items,
                ..
            } = item_trait;
            let ref self_ty: Box<Type> = parse_quote! {
                // use a less mangled name for more readable lifetime bounds
                // erros
                /* __require_unsafe */ __Self
            };
            let ref mut impl_trait_param: TypeParam = parse_quote! {
                #self_ty : ?Sized
            };
            impl_trait_param.bounds.push({
                let (
                    _impl_generics,
                    ty_generics,
                    _where_clause,
                ) = generics.split_for_impl();
                let trait_name = &item_trait.ident;
                parse_quote! {
                    #trait_name #ty_generics
                }
            });
            let supertraits = Vec::from_iter(
                item_trait
                    .supertraits
                    .iter()
                    .cloned()
                    .chain(iter::once({
                        let (
                            _impl_generics,
                            ty_generics,
                            _where_clause,
                        ) = generics.split_for_impl();
                        let trait_name = &item_trait.ident;
                        parse_quote! {
                            #trait_name #ty_generics
                        }
                    }))
            );

            unwrap!(
                items
                    .iter_mut()
                    .filter_map(|trait_item| match *trait_item {
                        | TraitItem::Method(TraitItemMethod {
                            ref attrs,
                            ref mut sig,
                            default: Some(ref mut block),
                            ..
                        }) => {
                            Some(FunctionLike {
                                attrs,
                                sig,
                                block,
                            })
                        },

                        | _ => None,
                    })
                    .try_for_each(|method| {
                        edit_method_in_place(
                            method,
                            Some((
                                generics,
                                self_ty,
                                Some((
                                    Some(impl_trait_param),
                                    &supertraits,
                                )),
                            )),
                        )
                    })
            );
            TokenStream::from(quote! {
                #item_trait
            })
        },
    }
}

#[cfg(all(test, feature = "unit-tests"))]
mod tests;