// Copyright (C) 2023 Parity Technologies (UK) Ltd. (admin@parity.io)
// This file is a part of the scale-encode crate.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//         http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use proc_macro2::TokenStream as TokenStream2;
use quote::{format_ident, quote};
use syn::{parse_macro_input, punctuated::Punctuated, DeriveInput};

const ATTR_NAME: &str = "encode_as_type";

/// The `EncodeAsType` derive macro can be used to implement `EncodeAsType`
/// on structs and enums whose fields all implement `EncodeAsType`.
///
/// # Examples
///
/// This can be applied to structs and enums:
///
/// ```rust
/// use scale_encode::EncodeAsType;
///
/// #[derive(EncodeAsType)]
/// struct Foo(String);
///
/// #[derive(EncodeAsType)]
/// struct Bar {
///     a: u64,
///     b: bool
/// }
///
/// #[derive(EncodeAsType)]
/// enum Wibble<T> {
///     A(usize, bool, T),
///     B { value: String },
///     C
/// }
/// ```
///
/// If you aren't directly depending on `scale_encode`, you must tell the macro what the path
/// to it is so that it knows how to generate the relevant impls:
///
/// ```rust
/// # use scale_encode as alt_path;
/// use alt_path::EncodeAsType;
///
/// #[derive(EncodeAsType)]
/// #[encode_as_type(crate_path = "alt_path")]
/// struct Foo<T> {
///    a: u64,
///    b: T
/// }
/// ```
///
/// If you use generics, the macro will assume that each of them also implements `EncodeAsType`.
/// This can be overridden when it's not the case (the compiler will ensure that you can't go wrong here):
///
/// ```rust
/// use scale_encode::EncodeAsType;
///
/// #[derive(EncodeAsType)]
/// #[encode_as_type(trait_bounds = "")]
/// struct Foo<T> {
///    a: u64,
///    b: bool,
///    c: std::marker::PhantomData<T>
/// }
/// ```
///
/// # Attributes
///
/// - `#[encode_as_type(crate_path = "::path::to::scale_encode")]`:
///   By default, the macro expects `scale_encode` to be a top level dependency,
///   available as `::scale_encode`. If this is not the case, you can provide the
///   crate path here.
/// - `#[encode_as_type(trait_bounds = "T: Foo, U::Input: EncodeAsType")]`:
///   By default, for each generate type parameter, the macro will add trait bounds such
///   that these type parameters must implement `EncodeAsType` too. You can override this
///   behaviour and provide your own trait bounds instead using this option.
#[proc_macro_derive(EncodeAsType, attributes(encode_as_type))]
pub fn derive_macro(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let input = parse_macro_input!(input as DeriveInput);

    // parse top level attrs.
    let attrs = match TopLevelAttrs::parse(&input.attrs) {
        Ok(attrs) => attrs,
        Err(e) => return e.write_errors().into(),
    };

    derive_with_attrs(attrs, input).into()
}

fn derive_with_attrs(attrs: TopLevelAttrs, input: DeriveInput) -> TokenStream2 {
    // what type is the derive macro declared on?
    match &input.data {
        syn::Data::Enum(details) => generate_enum_impl(attrs, &input, details),
        syn::Data::Struct(details) => generate_struct_impl(attrs, &input, details),
        syn::Data::Union(_) => syn::Error::new(
            input.ident.span(),
            "Unions are not supported by the EncodeAsType macro",
        )
        .into_compile_error(),
    }
}

fn generate_enum_impl(
    attrs: TopLevelAttrs,
    input: &DeriveInput,
    details: &syn::DataEnum,
) -> TokenStream2 {
    let path_to_scale_encode = &attrs.crate_path;
    let path_to_type: syn::Path = input.ident.clone().into();
    let (impl_generics, ty_generics, where_clause) = handle_generics(&attrs, &input.generics);

    // For each variant we want to spit out a match arm.
    let match_arms = details.variants.iter().map(|variant| {
        let variant_name = &variant.ident;
        let variant_name_str = variant_name.to_string();

        let (matcher, composite) =
            fields_to_matcher_and_composite(path_to_scale_encode, &variant.fields);
        quote!(
            Self::#variant_name #matcher => {
                #path_to_scale_encode::Variant { name: #variant_name_str, fields: #composite }
                    .encode_as_type_to(
                        __encode_as_type_type_id,
                        __encode_as_type_types,
                        __encode_as_type_out
                    )
            }
        )
    });

    quote!(
        impl #impl_generics #path_to_scale_encode::EncodeAsType for #path_to_type #ty_generics #where_clause {
            fn encode_as_type_to(
                &self,
                // long variable names to prevent conflict with struct field names:
                __encode_as_type_type_id: u32,
                __encode_as_type_types: &#path_to_scale_encode::PortableRegistry,
                __encode_as_type_out: &mut Vec<u8>
            ) -> Result<(), #path_to_scale_encode::Error> {
                match self {
                    #( #match_arms, )*
                    // This will never be encountered, but in case the enum has no variants
                    // the compiler will still want something to be spat out here:
                    _ => unreachable!()
                }
            }
        }
    )
}

fn generate_struct_impl(
    attrs: TopLevelAttrs,
    input: &DeriveInput,
    details: &syn::DataStruct,
) -> TokenStream2 {
    let path_to_scale_encode = &attrs.crate_path;
    let path_to_type: syn::Path = input.ident.clone().into();
    let (impl_generics, ty_generics, where_clause) = handle_generics(&attrs, &input.generics);

    let (matcher, composite) =
        fields_to_matcher_and_composite(path_to_scale_encode, &details.fields);

    quote!(
        impl #impl_generics #path_to_scale_encode::EncodeAsType for #path_to_type #ty_generics #where_clause {
            fn encode_as_type_to(
                &self,
                // long variable names to prevent conflict with struct field names:
                __encode_as_type_type_id: u32,
                __encode_as_type_types: &#path_to_scale_encode::PortableRegistry,
                __encode_as_type_out: &mut Vec<u8>
            ) -> Result<(), #path_to_scale_encode::Error> {
                let #path_to_type #matcher = self;
                #composite.encode_as_type_to(
                    __encode_as_type_type_id,
                    __encode_as_type_types,
                    __encode_as_type_out
                )
            }
        }
        impl #impl_generics #path_to_scale_encode::EncodeAsFields for #path_to_type #ty_generics #where_clause {
            fn encode_as_fields_to<'__encode_as_field_lt, I: #path_to_scale_encode::FieldIter<'__encode_as_field_lt>>(
                &self,
                // long variable names to prevent conflict with struct field names:
                __encode_as_type_fields: I,
                __encode_as_type_types: &#path_to_scale_encode::PortableRegistry,
                __encode_as_type_out: &mut Vec<u8>
            ) -> Result<(), #path_to_scale_encode::Error> {
                let #path_to_type #matcher = self;
                #composite.encode_as_fields_to(
                    __encode_as_type_fields,
                    __encode_as_type_types,
                    __encode_as_type_out
                )
            }
        }
    )
}

fn handle_generics<'a>(
    attrs: &TopLevelAttrs,
    generics: &'a syn::Generics,
) -> (
    syn::ImplGenerics<'a>,
    syn::TypeGenerics<'a>,
    syn::WhereClause,
) {
    let path_to_crate = &attrs.crate_path;
    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();

    let mut where_clause = where_clause.cloned().unwrap_or(syn::parse_quote!(where));

    if let Some(where_predicates) = &attrs.trait_bounds {
        // if custom trait bounds are given, append those to the where clause.
        where_clause.predicates.extend(where_predicates.clone());
    } else {
        // else, append our default EncodeAsType bounds to the where clause.
        for param in generics.type_params() {
            let ty = &param.ident;
            where_clause
                .predicates
                .push(syn::parse_quote!(#ty: #path_to_crate::EncodeAsType))
        }
    }

    (impl_generics, ty_generics, where_clause)
}

fn fields_to_matcher_and_composite(
    path_to_scale_encode: &syn::Path,
    fields: &syn::Fields,
) -> (TokenStream2, TokenStream2) {
    match fields {
        syn::Fields::Named(fields) => {
            let match_body = fields.named.iter().map(|f| {
                let field_name = &f.ident;
                quote!(#field_name)
            });
            let tuple_body = fields.named.iter().map(|f| {
                let field_name_str = f.ident.as_ref().unwrap().to_string();
                let field_name = &f.ident;
                quote!((Some(#field_name_str), #field_name as &dyn #path_to_scale_encode::EncodeAsType))
            });

            (
                quote!({#( #match_body ),*}),
                quote!(#path_to_scale_encode::Composite([#( #tuple_body ),*].into_iter())),
            )
        }
        syn::Fields::Unnamed(fields) => {
            let field_idents: Vec<syn::Ident> = fields
                .unnamed
                .iter()
                .enumerate()
                .map(|(idx, _)| format_ident!("_{idx}"))
                .collect();
            let match_body = field_idents.iter().map(|i| quote!(#i));
            let tuple_body = field_idents
                .iter()
                .map(|i| quote!((None as Option<&'static str>, #i as &dyn #path_to_scale_encode::EncodeAsType)));

            (
                quote!((#( #match_body ),*)),
                quote!(#path_to_scale_encode::Composite([#( #tuple_body ),*].into_iter())),
            )
        }
        syn::Fields::Unit => (
            quote!(),
            quote!(#path_to_scale_encode::Composite(([] as [(Option<&'static str>, &dyn #path_to_scale_encode::EncodeAsType);0]).into_iter())),
        ),
    }
}

struct TopLevelAttrs {
    // path to the scale_encode crate, in case it's not a top level dependency.
    crate_path: syn::Path,
    // allow custom trait bounds to be used instead of the defaults.
    trait_bounds: Option<Punctuated<syn::WherePredicate, syn::Token!(,)>>,
}

impl TopLevelAttrs {
    fn parse(attrs: &[syn::Attribute]) -> darling::Result<Self> {
        use darling::FromMeta;

        #[derive(FromMeta)]
        struct TopLevelAttrsInner {
            #[darling(default)]
            crate_path: Option<syn::Path>,
            #[darling(default)]
            trait_bounds: Option<Punctuated<syn::WherePredicate, syn::Token!(,)>>,
        }

        let mut res = TopLevelAttrs {
            crate_path: syn::parse_quote!(::scale_encode),
            trait_bounds: None,
        };

        // look at each top level attr. parse any for encode_as_type.
        for attr in attrs {
            if !attr.path.is_ident(ATTR_NAME) {
                continue;
            }
            let meta = attr.parse_meta()?;
            let parsed_attrs = TopLevelAttrsInner::from_meta(&meta)?;

            res.trait_bounds = parsed_attrs.trait_bounds;
            if let Some(crate_path) = parsed_attrs.crate_path {
                res.crate_path = crate_path;
            }
        }

        Ok(res)
    }
}