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)| 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;

                    None
                }
                else {
                    prev_token_was_prefix = false;

                    Some(TokenTree::Punct(punct))
                }
            }
            TokenTree::Ident(ident) if prev_token_was_prefix && &ident == target_ident => {
                prev_token_was_prefix = false;

                Some(substitution.clone())
            }
            TokenTree::Ident(_) if prev_token_was_prefix => {
                prev_token_was_prefix = false;

                Some(substitution.clone())
            }
            TokenTree::Group(group) => {
                prev_token_was_prefix = false;

                Some(TokenTree::Group(Group::new(
                    group.delimiter(),
                    recurse_replace_ident(group.stream(), target_ident, substitution, prefix),
                )))
            }
            tt => {
                prev_token_was_prefix = false;

                Some(tt)
            }
        })
        .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);
            }
            .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 #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"))),
            )
        });
    }
}