use crate::{
    error::Error,
    utils::{Cat3, Comma, CommaPunctuated, IdentAssignment},
};
use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use syn::{
    DeriveInput, Error as SynError, Path, Type,
    parse::{Parse, ParseStream},
};

pub struct TypeAssoc {
    trait_path: Path,
    assoc_type_token_streams: Vec<TokenStream2>,
}

impl TypeAssoc {
    const TYPE_ASSOC_ATTRIBUTE_NAME: &str = "type_assoc";
    const TRAIT_PATH_KEY: &str = "impl_trait";

    fn assoc_type_token_stream(Cat3(assoc_type_ident, _comma, assoc_type_type): IdentAssignment<Type>) -> TokenStream2 {
        quote::quote! {
            type #assoc_type_ident = #assoc_type_type;
        }
    }

    fn from_derive_input(input: &DeriveInput) -> Result<Self, SynError> {
        for attribute in &input.attrs {
            if attribute.path().is_ident(Self::TYPE_ASSOC_ATTRIBUTE_NAME) {
                return attribute.parse_args();
            }
        }

        Err(Error::missing_expected_attribute(
            input,
            Self::TYPE_ASSOC_ATTRIBUTE_NAME,
        ))
    }

    pub fn derive_impl(input: &DeriveInput) -> Result<TokenStream2, SynError> {
        let input_ident = &input.ident;
        let Self {
            trait_path,
            assoc_type_token_streams,
        } = Self::from_derive_input(input)?;
        let (impl_generics, input_generics, input_where_clause) = input.generics.split_for_impl();
        let impl_block_token_stream = quote::quote! {
            impl #impl_generics #trait_path for #input_ident #input_generics #input_where_clause {
                #(#assoc_type_token_streams)*
            }
        };

        Ok(impl_block_token_stream)
    }

    pub fn derive(input_token_stream: TokenStream) -> TokenStream {
        let input = syn::parse_macro_input!(input_token_stream);

        Self::derive_impl(&input)
            .unwrap_or_else(SynError::into_compile_error)
            .into()
    }
}

impl Parse for TypeAssoc {
    // NOTE: would ideally use [Cat3<IdentAssignment<Path>, Comma, CommaPunctuated<IdentAssignment<Type>>>]
    // but [Punctuated] does not implement [Parse]
    fn parse(parse_stream: ParseStream) -> Result<Self, SynError> {
        let (trait_path_key, _equals, trait_path) = parse_stream.parse::<IdentAssignment<Path>>()?.into_tuple();

        if trait_path_key != Self::TRAIT_PATH_KEY {
            return Err(Error::unexpected_value(&trait_path_key, Self::TRAIT_PATH_KEY));
        }

        parse_stream.parse::<Comma>()?;

        let assoc_type_assignments = CommaPunctuated::<IdentAssignment<Type>>::parse_terminated(parse_stream)?;
        let assoc_type_token_streams = assoc_type_assignments
            .into_iter()
            .map(Self::assoc_type_token_stream)
            .collect();
        let type_assoc = Self {
            trait_path,
            assoc_type_token_streams,
        };

        Ok(type_assoc)
    }
}
