extern crate proc_macro;

use std::collections::HashMap;

use proc_macro2::TokenStream;
use quote::quote;
use syn::parse::Parse;
use syn::{Error, Result};

use crate::{filler, Fill};

/// One-line wrapper that declares a filler macro.
///
/// # Example
/// ```
/// # extern crate proc_macro;
/// #
/// portrait_framework::proc_macro_filler!(foo, Generator);
/// struct Generator(portrait_framework::NoArgs);
/// impl portrait_framework::Generate for Generator {
///     fn generate_const(
///         &mut self,
///         context: portrait_framework::Context,
///         item: &syn::TraitItemConst,
///     ) -> syn::Result<syn::ImplItemConst> {
///         todo!()
///     }
///     fn generate_fn(
///         &mut self,
///         context: portrait_framework::Context,
///         item: &syn::TraitItemFn,
///     ) -> syn::Result<syn::ImplItemFn> {
///         todo!()
///     }
///     fn generate_type(
///         &mut self,
///         context: portrait_framework::Context,
///         item: &syn::TraitItemType,
///     ) -> syn::Result<syn::ImplItemType> {
///         todo!()
///     }
/// }
/// ```
///
/// This declares a filler macro called `foo`,
/// where each missing item is generated by calling the corresponding funciton.
#[macro_export]
macro_rules! proc_macro_filler {
    ($ident:ident, $generator:path) => {
        pub fn $ident(input: ::proc_macro::TokenStream) -> ::proc_macro::TokenStream {
            portrait_framework::completer_filler(input, $generator)
        }
    };
}

/// Shorthand from [`fn@filler`] to [`complete`] ([`proc_macro`] version).
pub fn completer_filler<ArgsT: Parse, GeneratorT: Generate>(
    input: proc_macro::TokenStream,
    ctor: fn(ArgsT) -> GeneratorT,
) -> proc_macro::TokenStream {
    completer_filler2(input.into(), ctor).unwrap_or_else(syn::Error::into_compile_error).into()
}

/// Shorthand from [`fn@filler`] to [`complete`] ([`proc_macro2`] version).
pub fn completer_filler2<ArgsT: Parse, GeneratorT: Generate>(
    input: TokenStream,
    ctor: fn(ArgsT) -> GeneratorT,
) -> Result<TokenStream> {
    struct Filler<GenerateT, ArgsT>(fn(ArgsT) -> GenerateT);

    impl<GenerateT: Generate, ArgsT: Parse> Fill for Filler<GenerateT, ArgsT> {
        type Args = ArgsT;

        fn fill(
            self,
            portrait: &[syn::TraitItem],
            args: Self::Args,
            item_impl: &syn::ItemImpl,
        ) -> Result<TokenStream> {
            let tokens = complete(portrait, item_impl, self.0(args))?;
            Ok(quote!(#tokens))
        }
    }

    filler(input, Filler(ctor))
}

/// Invokes the generator on each unimplemented item
/// and returns a clone of `impl_block` with the generated items.
pub fn complete(
    trait_items: &[syn::TraitItem],
    impl_block: &syn::ItemImpl,
    mut generator: impl Generate,
) -> syn::Result<syn::ItemImpl> {
    let mut output = impl_block.clone();

    let ctx = Context { all_trait_items: trait_items, impl_block };

    let items = subtract_items(trait_items, impl_block)?;
    for trait_item in items.consts.values() {
        let impl_item = generator.generate_const(Context { ..ctx }, trait_item)?;
        output.items.push(syn::ImplItem::Const(impl_item));
    }
    for trait_item in items.fns.values() {
        let impl_item = generator.generate_fn(Context { ..ctx }, trait_item)?;
        output.items.push(syn::ImplItem::Fn(impl_item));
    }
    for trait_item in items.types.values() {
        let impl_item = generator.generate_type(Context { ..ctx }, trait_item)?;
        output.items.push(syn::ImplItem::Type(impl_item));
    }

    Ok(output)
}

/// Available context parameters passed to generators.
#[non_exhaustive]
pub struct Context<'t> {
    /// All known trait items in the portrait.
    pub all_trait_items: &'t [syn::TraitItem],
    /// The input impl block.
    pub impl_block:      &'t syn::ItemImpl,
}

/// Generates missing items.
pub trait Generate {
    /// Implements an associated constant.
    fn generate_const(
        &mut self,
        ctx: Context,
        item: &syn::TraitItemConst,
    ) -> Result<syn::ImplItemConst>;

    /// Implements an associated function.
    fn generate_fn(&mut self, ctx: Context, item: &syn::TraitItemFn) -> Result<syn::ImplItemFn>;

    /// Implements an associated type.
    fn generate_type(
        &mut self,
        ctx: Context,
        item: &syn::TraitItemType,
    ) -> Result<syn::ImplItemType>;
}

/// Shorthand for `TraitItemMap::new().minus(ImplItemMap::new())`.
pub fn subtract_items<'t>(
    trait_items: &'t [syn::TraitItem],
    impl_block: &'t syn::ItemImpl,
) -> syn::Result<TraitItemMap<'t>> {
    let mut items = TraitItemMap::new(trait_items);
    items.minus(&ImplItemMap::new(impl_block))?;
    Ok(items)
}

/// Indexes items in a trait by namespaced identifier.
#[derive(Default)]
pub struct TraitItemMap<'t> {
    /// Associated constants in the trait.
    pub consts: HashMap<syn::Ident, &'t syn::TraitItemConst>,
    /// Associated functions in the trait.
    pub fns:    HashMap<syn::Ident, &'t syn::TraitItemFn>,
    /// Associated types in the trait.
    pub types:  HashMap<syn::Ident, &'t syn::TraitItemType>,
}

impl<'t> TraitItemMap<'t> {
    /// Constructs the trait item index from a slice of trait items.
    pub fn new(trait_items: &'t [syn::TraitItem]) -> Self {
        let mut map = Self::default();
        for item in trait_items {
            match item {
                syn::TraitItem::Const(item) => {
                    map.consts.insert(item.ident.clone(), item);
                }
                syn::TraitItem::Fn(item) => {
                    map.fns.insert(item.sig.ident.clone(), item);
                }
                syn::TraitItem::Type(item) => {
                    map.types.insert(item.ident.clone(), item);
                }
                _ => {}
            }
        }
        map
    }

    /// Removes the items found in the impl, leaving only unimplemented items.
    pub fn minus(&mut self, impl_items: &ImplItemMap) -> Result<()> {
        for (ident, impl_item) in &impl_items.consts {
            if self.consts.remove(ident).is_none() {
                return Err(Error::new_spanned(
                    impl_item,
                    "no associated constant called {ident} in trait",
                ));
            }
        }

        for (ident, impl_item) in &impl_items.fns {
            if self.fns.remove(ident).is_none() {
                return Err(Error::new_spanned(
                    impl_item,
                    "no associated function called {ident} in trait",
                ));
            }
        }

        for (ident, impl_item) in &impl_items.types {
            if self.types.remove(ident).is_none() {
                return Err(Error::new_spanned(
                    impl_item,
                    "no associated type called {ident} in trait",
                ));
            }
        }

        Ok(())
    }
}

/// Indexes items in an impl block by namespaced identifier.
#[derive(Default)]
pub struct ImplItemMap<'t> {
    /// Associated constants in the implementation.
    pub consts: HashMap<syn::Ident, &'t syn::ImplItemConst>,
    /// Associated functions in the implementation.
    pub fns:    HashMap<syn::Ident, &'t syn::ImplItemFn>,
    /// Associated types in the implementation.
    pub types:  HashMap<syn::Ident, &'t syn::ImplItemType>,
}

impl<'t> ImplItemMap<'t> {
    /// Constructs the impl item index from an impl block.
    pub fn new(impl_block: &'t syn::ItemImpl) -> Self {
        let mut map = Self::default();
        for item in &impl_block.items {
            match item {
                syn::ImplItem::Const(item) => {
                    map.consts.insert(item.ident.clone(), item);
                }
                syn::ImplItem::Fn(item) => {
                    map.fns.insert(item.sig.ident.clone(), item);
                }
                syn::ImplItem::Type(item) => {
                    map.types.insert(item.ident.clone(), item);
                }
                _ => {}
            }
        }
        map
    }
}