#![doc = include_str!("../README.md")]
use std::collections::HashMap;

use proc_macro::TokenStream;
use proc_macro2::Group;
use proc_macro2::Ident;
use quote::format_ident;
use quote::quote;
use quote::quote_spanned;
use quote::ToTokens;
use syn::parse_macro_input;
use syn::punctuated::Punctuated;
use syn::token::Brace;
use syn::token::Gt;
use syn::token::Lt;
use syn::token::Trait;
use syn::AngleBracketedGenericArguments;
use syn::Block;
use syn::Expr;
use syn::GenericArgument;
use syn::GenericParam;
use syn::ImplItemMethod;
use syn::ItemTrait;
use syn::Path;
use syn::PathArguments;
use syn::PathSegment;
use syn::Signature;
use syn::Stmt;
use syn::TraitBound;
use syn::TraitItem;
use syn::TraitItemMethod;
use syn::Type;
use syn::TypeParam;
use syn::TypeParamBound;
use syn::TypePath;
use syn::Visibility;
use syn_utils::TypeMatcher;

use crate::parse_assoc_type::parse_assoc_type;
use crate::parse_assoc_type::AssocTypeParseError;
use crate::parse_trait_sig::parse_trait_signature;
use crate::parse_trait_sig::MethodParseError;
use crate::parse_trait_sig::SignatureChanges;
use crate::parse_trait_sig::TypeTransform;
mod parse_assoc_type;
mod parse_trait_sig;
mod syn_utils;

macro_rules! abort {
    ($span:expr, $message:literal) => {
        quote_spanned! {$span => compile_error!($message);}.into()
    };
}

#[proc_macro_attribute]
pub fn dynamize(_attr: TokenStream, input: TokenStream) -> TokenStream {
    let mut original_trait = parse_macro_input!(input as ItemTrait);
    assert!(original_trait.auto_token.is_none());

    let original_trait_name = original_trait.ident.clone();

    let mut objectifiable_methods: Vec<(Signature, SignatureChanges)> = Vec::new();

    let mut assoc_type_conversions: HashMap<&Ident, &Type> = HashMap::new();

    for item in &original_trait.items {
        if let TraitItem::Type(assoc_type) = item {
            match parse_assoc_type(assoc_type) {
                Err((_, AssocTypeParseError::NoIntoBound)) => continue,
                Err((span, AssocTypeParseError::AssocTypeInBound)) => {
                    return abort!(span, "dynamize does not support associated types here")
                }
                Err((span, AssocTypeParseError::GenericAssociatedType)) => {
                    return abort!(
                        span,
                        "dynamize does not (yet?) support generic associated types"
                    )
                }
                Ok((ident, type_)) => {
                    assoc_type_conversions.insert(ident, type_);
                }
            }
        }
    }

    for item in &original_trait.items {
        if let TraitItem::Method(method) = item {
            let mut signature = method.sig.clone();
            match parse_trait_signature(&mut signature, &assoc_type_conversions) {
                Err((_, MethodParseError::NonDispatchableMethod)) => continue,
                Err((span, MethodParseError::UnconvertibleAssocType)) => {
                    return abort!(
                        span,
                        "associated type is either undefined or doesn't have an Into bound"
                    )
                }
                Err((span, MethodParseError::AssocTypeInInputs)) => {
                    return abort!(
                        span,
                        "dynamize does not support associated types in parameter types"
                    )
                }
                Err((
                    span,
                    MethodParseError::AssocTypeInUnsupportedReturnType
                    | MethodParseError::UnconvertibleAssocTypeInFnInput,
                )) => return abort!(span, "dynamize does not know how to convert this type"),
                Err((span, MethodParseError::UnconvertibleAssocTypeInTraitBound)) => {
                    return abort!(span, "dynamize does not support associated types here")
                }
                Err((span, MethodParseError::ImplTraitInInputs)) => {
                    return abort!(
                        span,
                        "dynamize does not support impl here, change it to a method generic"
                    )
                }
                Ok(parsed_method) => objectifiable_methods.push((signature, parsed_method)),
            };
        }
    }

    let mut method_impls: Vec<ImplItemMethod> = Vec::new();

    let mut blanket_impl_attrs = Vec::new();
    let mut dyn_trait_attrs = Vec::new();

    // FUTURE: use Vec::drain_filter once it's stable
    let mut i = 0;
    while i < original_trait.attrs.len() {
        if original_trait.attrs[i].path.is_ident("blanket_impl_attr") {
            let attr = original_trait.attrs.remove(i);
            let group: Group = match syn::parse2(attr.tokens) {
                Ok(g) => g,
                Err(err) => {
                    return abort!(
                        err.span(),
                        "expected parenthesis: #[blanket_impl_attr(...)]"
                    )
                }
            };
            let tokens = group.stream();
            blanket_impl_attrs.push(quote! {#[#tokens]});
        } else if original_trait.attrs[i].path.is_ident("dyn_trait_attr") {
            let attr = original_trait.attrs.remove(i);
            let group: Group = match syn::parse2(attr.tokens) {
                Ok(g) => g,
                Err(err) => {
                    return abort!(err.span(), "expected parenthesis: #[dyn_trait_attr(...)]")
                }
            };
            let tokens = group.stream();
            dyn_trait_attrs.push(quote! {#[#tokens]});
        } else {
            i += 1;
        }
    }

    let mut dyn_trait = ItemTrait {
        ident: format_ident!("Dyn{}", original_trait.ident),

        attrs: Vec::new(),
        vis: original_trait.vis.clone(),
        unsafety: original_trait.unsafety,
        auto_token: None,
        trait_token: Trait::default(),
        generics: original_trait.generics.clone(),
        colon_token: None,
        supertraits: Punctuated::new(),
        brace_token: Brace::default(),
        items: Vec::new(),
    };

    for (signature, parsed_method) in objectifiable_methods {
        let mut new_method = TraitItemMethod {
            attrs: Vec::new(),
            sig: signature,
            default: None,
            semi_token: None,
        };

        let fun_name = &new_method.sig.ident;

        let args = new_method
            .sig
            .inputs
            .iter()
            .enumerate()
            .map(|(idx, arg)| match arg {
                syn::FnArg::Receiver(_) => quote! {self},
                syn::FnArg::Typed(pat_type) => match pat_type.pat.as_ref() {
                    syn::Pat::Ident(ident) => match &parsed_method.inputs[idx] {
                        None => ident.ident.to_token_stream(),
                        Some(transforms) => {
                            let args = (0..transforms.len()).map(|i| format_ident!("a{}", i));
                            let mut calls: Vec<_> =
                                args.clone().map(|i| i.into_token_stream()).collect();
                            for i in 0..calls.len() {
                                transforms[i].append_conversion(&mut calls[i]);
                            }
                            let move_opt = new_method.sig.asyncness.map(|_| quote! {move});
                            quote!(#move_opt |#(#args),*| #ident(#(#calls),*))
                        }
                    },
                    _other => {
                        panic!("unexpected");
                    }
                },
            });

        // in order for a trait to be object-safe its methods may not have
        // generics so we convert method generics into trait generics
        if new_method
            .sig
            .generics
            .params
            .iter()
            .any(|p| matches!(p, GenericParam::Type(_)))
        {
            // syn::punctuated::Punctuated doesn't have a remove(index)
            // method so we firstly move all elements to a vector
            let mut params = Vec::new();
            while let Some(generic_param) = new_method.sig.generics.params.pop() {
                params.push(generic_param.into_value());
            }

            // FUTURE: use Vec::drain_filter once it's stable
            let mut i = 0;
            while i < params.len() {
                if matches!(params[i], GenericParam::Type(_)) {
                    dyn_trait.generics.params.push(params.remove(i));
                } else {
                    i += 1;
                }
            }

            new_method.sig.generics.params.extend(params);

            if dyn_trait.generics.lt_token.is_none() {
                dyn_trait.generics.lt_token = Some(Lt::default());
                dyn_trait.generics.gt_token = Some(Gt::default());
            }
        }

        let mut expr = quote!(#original_trait_name::#fun_name(#(#args),*));
        if new_method.sig.asyncness.is_some() {
            expr.extend(quote! {.await})
        }
        parsed_method.return_type.append_conversion(&mut expr);

        method_impls.push(ImplItemMethod {
            attrs: Vec::new(),
            vis: Visibility::Inherited,
            defaultness: None,
            sig: new_method.sig.clone(),
            block: Block {
                brace_token: Brace::default(),
                stmts: vec![Stmt::Expr(Expr::Verbatim(expr))],
            },
        });
        dyn_trait.items.push(new_method.into());
    }

    let blanket_impl = generate_blanket_impl(&dyn_trait, &original_trait, &method_impls);

    let expanded = quote! {
        #original_trait

        #(#dyn_trait_attrs)*
        #dyn_trait

        #(#blanket_impl_attrs)*
        #blanket_impl
    };
    TokenStream::from(expanded)
}

fn generate_blanket_impl(
    dyn_trait: &ItemTrait,
    original_trait: &ItemTrait,
    method_impls: &[ImplItemMethod],
) -> proc_macro2::TokenStream {
    let mut blanket_generics = dyn_trait.generics.clone();
    let some_ident = format_ident!("__to_be_dynamized");
    blanket_generics.params.push(GenericParam::Type(TypeParam {
        attrs: Vec::new(),
        ident: some_ident.clone(),
        colon_token: None,
        bounds: std::iter::once(TypeParamBound::Trait(TraitBound {
            paren_token: None,
            modifier: syn::TraitBoundModifier::None,
            lifetimes: None,
            path: Path {
                leading_colon: None,
                segments: std::iter::once(path_segment_for_trait(original_trait)).collect(),
            },
        }))
        .collect(),
        eq_token: None,
        default: None,
    }));
    let (_, type_gen, _where) = dyn_trait.generics.split_for_impl();
    let dyn_trait_name = &dyn_trait.ident;
    quote! {
        impl #blanket_generics #dyn_trait_name #type_gen for #some_ident {
            #(#method_impls)*
        }
    }
}

struct AssocTypeMatcher;
impl TypeMatcher<Path> for AssocTypeMatcher {
    fn match_path<'a>(&self, path: &'a Path) -> Option<&'a Path> {
        if path.segments.first().unwrap().ident == "Self" {
            return Some(path);
        }
        None
    }
}

impl<TypeMatch, PathMatch, T> TypeMatcher<T> for (TypeMatch, PathMatch)
where
    TypeMatch: Fn(&Type) -> Option<&T>,
    PathMatch: Fn(&Path) -> Option<&T>,
{
    fn match_type<'a>(&self, t: &'a Type) -> Option<&'a T> {
        self.0(t)
    }

    fn match_path<'a>(&self, path: &'a Path) -> Option<&'a T> {
        self.1(path)
    }
}

impl TypeTransform {
    fn append_conversion(&self, stream: &mut proc_macro2::TokenStream) {
        match self {
            TypeTransform::Into => stream.extend(quote! {.into()}),
            TypeTransform::Map(opt) => {
                let mut inner = quote!(x);
                opt.append_conversion(&mut inner);
                stream.extend(quote! {.map(|x| #inner)})
            }
            TypeTransform::Result(ok, err) => {
                if !matches!(ok.as_ref(), TypeTransform::NoOp) {
                    let mut inner = quote!(x);
                    ok.append_conversion(&mut inner);
                    stream.extend(quote! {.map(|x| #inner)})
                }
                if !matches!(err.as_ref(), TypeTransform::NoOp) {
                    let mut inner = quote!(x);
                    err.append_conversion(&mut inner);
                    stream.extend(quote! {.map_err(|x| #inner)})
                }
            }
            _other => {}
        }
    }
}

/// Just a convenience trait for us to avoid match/if-let blocks everywhere.
trait As<T> {
    fn get_as(&self) -> Option<&T>;
    fn get_as_mut(&mut self) -> Option<&mut T>;
}

impl As<AngleBracketedGenericArguments> for PathArguments {
    fn get_as(&self) -> Option<&AngleBracketedGenericArguments> {
        match self {
            PathArguments::AngleBracketed(args) => Some(args),
            _other => None,
        }
    }
    fn get_as_mut(&mut self) -> Option<&mut AngleBracketedGenericArguments> {
        match self {
            PathArguments::AngleBracketed(args) => Some(args),
            _other => None,
        }
    }
}

impl As<Type> for GenericArgument {
    fn get_as(&self) -> Option<&Type> {
        match self {
            GenericArgument::Type(typearg) => Some(typearg),
            _other => None,
        }
    }
    fn get_as_mut(&mut self) -> Option<&mut Type> {
        match self {
            GenericArgument::Type(typearg) => Some(typearg),
            _other => None,
        }
    }
}

fn path_segment_for_trait(sometrait: &ItemTrait) -> PathSegment {
    PathSegment {
        ident: sometrait.ident.clone(),
        arguments: match sometrait.generics.params.is_empty() {
            true => PathArguments::None,
            false => PathArguments::AngleBracketed(AngleBracketedGenericArguments {
                colon2_token: None,
                lt_token: Lt::default(),
                args: sometrait
                    .generics
                    .params
                    .iter()
                    .map(|param| match param {
                        GenericParam::Type(type_param) => {
                            GenericArgument::Type(Type::Path(TypePath {
                                path: type_param.ident.clone().into(),
                                qself: None,
                            }))
                        }
                        GenericParam::Lifetime(lifetime_def) => {
                            GenericArgument::Lifetime(lifetime_def.lifetime.clone())
                        }
                        GenericParam::Const(_) => todo!("const generic param not supported"),
                    })
                    .collect(),
                gt_token: Gt::default(),
            }),
        },
    }
}