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

use crate::iter::IteratorExt;

#[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,
    prefix: char,
) -> TokenStream
{
    let mut prev_token_was_prefix = false;

    token_stream
        .into_iter()
        .windows()
        .filter_map(|(token_tree, next_token_tree)| {
            let output = match token_tree {
                TokenTree::Punct(punct) if punct.as_char() == prefix => {
                    if matches!(
                        next_token_tree,
                        Some(TokenTree::Ident(ident)) if &ident == target_ident
                    ) {
                        prev_token_was_prefix = true;

                        return None;
                    }

                    Some(TokenTree::Punct(punct))
                }
                TokenTree::Ident(ident)
                    if prev_token_was_prefix && &ident == target_ident =>
                {
                    Some(substitution.clone())
                }
                TokenTree::Group(group) => Some(TokenTree::Group(Group::new(
                    group.delimiter(),
                    recurse_replace_ident(
                        group.stream(),
                        target_ident,
                        substitution,
                        prefix,
                    ),
                ))),
                tt => Some(tt),
            };

            prev_token_was_prefix = false;

            output
        })
        .collect()
}

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

    use std::hint::black_box;

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

    use super::*;

    #[test]
    fn replace_ident_works()
    {
        let ht = Punct::new('#', Spacing::Alone);

        assert_eq!(
            quote! {
                let abc = #ht 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 = (#ht 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(|_| #ht xyz);

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

                fn #ht hello() {}
            }
            .to_string()
        );

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

                    foo
                }

                let xyz = "123";

                println!("Hello {}", #ht 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
                }

                let xyz = "123";

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

    #[bench]
    fn bench_replace_ident(bencher: &mut Bencher)
    {
        let ht = Punct::new('#', Spacing::Alone);

        let input_strem = quote! {
            let #ht foo = |abc| {
                let x = "foo";

                let y = a_foo;

                let #ht foo = "Hello";

                #ht foo.to_string()
            };
        };

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