use proc_macro2::{Group, Ident, TokenStream, TokenTree};

#[allow(clippy::module_name_repetitions)]
pub trait TokenStreamExt
{
    /// Recursively replaces all occurences of a identifier with a token tree.
    fn replace_ident(
        &self,
        target_ident: &Ident,
        substitution: &TokenTree,
    ) -> TokenStream;
}

impl TokenStreamExt for TokenStream
{
    fn replace_ident(&self, target_ident: &Ident, substitution: &TokenTree)
        -> TokenStream
    {
        recurse_replace_ident(self.clone(), target_ident, substitution)
    }
}

fn recurse_replace_ident(
    token_stream: TokenStream,
    target_ident: &Ident,
    substitution: &TokenTree,
) -> TokenStream
{
    token_stream
        .into_iter()
        .map(|token_tree| match token_tree {
            TokenTree::Ident(ident) if &ident == target_ident => substitution.clone(),
            TokenTree::Group(group) => TokenTree::Group(Group::new(
                group.delimiter(),
                recurse_replace_ident(group.stream(), target_ident, substitution),
            )),
            tt => tt,
        })
        .collect()
}

#[cfg(test)]
mod tests
{
    extern crate test;

    use std::hint::black_box;

    use proc_macro2::Span;
    use quote::{format_ident, quote};
    use test::Bencher;

    use super::*;

    #[test]
    fn replace_ident_works()
    {
        assert_eq!(
            quote! {
                let abc = xyz;
            }
            .replace_ident(
                &format_ident!("xyz"),
                &TokenTree::Ident(Ident::new("foo", Span::call_site())),
            )
            .to_string(),
            quote! {
                let abc = foo;
            }
            .to_string()
        );

        assert_eq!(
            quote! {
                let abc = (xyz, "123");
            }
            .replace_ident(
                &format_ident!("xyz"),
                &TokenTree::Ident(Ident::new("foo", Span::call_site())),
            )
            .to_string(),
            quote! {
                let abc = (foo, "123");
            }
            .to_string()
        );

        assert_eq!(
            quote! {
                let abc = (hello, "123").iter_map(|_| xyz);
            }
            .replace_ident(
                &format_ident!("xyz"),
                &TokenTree::Ident(Ident::new("foo", Span::call_site())),
            )
            .to_string(),
            quote! {
                let abc = (hello, "123").iter_map(|_| foo);
            }
            .to_string()
        );

        assert_eq!(
            quote! {
                fn xyz(bar: &Bar) {
                    let foo = "baz";

                    foo
                }

                println!("Hello {}", xyz());
            }
            .replace_ident(
                &format_ident!("xyz"),
                &TokenTree::Ident(Ident::new("foo", Span::call_site())),
            )
            .to_string(),
            quote! {
                fn foo(bar: &Bar) {
                    let foo = "baz";

                    foo
                }

                println!("Hello {}", foo());
            }
            .to_string()
        );
    }

    #[bench]
    fn bench_replace_ident(bencher: &mut Bencher)
    {
        let input_strem = quote! {
            let foo = |abc| {
                let x = "foo";

                let y = a_foo;

                let foo = "Hello";

                foo.to_string()
            };
        };

        bencher.iter(|| {
            input_strem.replace_ident(
                black_box(&format_ident!("foo")),
                black_box(&TokenTree::Ident(format_ident!("foobar"))),
            )
        });
    }
}