use proc_macro::{Delimiter, Group, TokenStream, TokenTree};
use proc_macro2::Span;
use quote::{quote, quote_spanned};
use syn::{
    parse::{Parse, ParseStream},
    parse_macro_input, parse_quote, visit,
    visit::Visit,
    visit_mut,
    visit_mut::VisitMut,
    Block, Expr, ExprAsync, ExprAwait, ExprYield, ItemFn, Pat, PatWild, Stmt, Token,
};

#[proc_macro]
pub fn generator(item: TokenStream) -> TokenStream {
    let (mut block, pat) = {
        let arg_version = item.clone();

        match syn::parse::<ArgGenerator>(arg_version) {
            Ok(gen) => (gen.block, gen.pattern),
            Err(_) => {
                let group = Group::new(Delimiter::Brace, item);
                let stream = TokenTree::Group(group).into();
                (
                    parse_macro_input!(stream as Block),
                    Pat::Wild(PatWild {
                        attrs: vec![],
                        underscore_token: <Token!(_)>::default(),
                    }),
                )
            }
        }
    };

    let mut visitor = YieldVisitor::default();
    visitor.visit_block(&block);

    if visitor.errors.len() > 0 {
        let errors = visitor.errors.into_iter().map(|(error, span)| {
            quote_spanned! { span =>
                compile_error!(#error);
            }
        });
        let out = quote! {
            {
                #(#errors)*
            }
        };
        return out.into();
    }

    let type_hint = if visitor.found_exprs > visitor.found_statement_exprs {
        quote! { _ }
    } else {
        quote! { () }
    };

    let mut visitor = BlockVisitor {};
    visitor.visit_block_mut(&mut block);

    let tokens = quote! {
        {
            use ::generate::{Generator, GeneratorState, __support};

            let (mut __resume, mut __yield) = __support::generator_mem::<#type_hint, _>();

            let __await_resume = __resume.clone();
            let __await_yield = __yield.clone();
            let __yield_awaiter = move |val| __support::yield_future(__await_resume.clone(), __await_yield.clone(), val);

            let build = move |#pat| async move {
                #block
            };

            __support::generator_for(__resume, __yield, build)
        }
    };

    tokens.into()
}

#[allow(unused)]
struct ArgGenerator {
    left_or: Token![|],
    pattern: Pat,
    right_or: Token![|],
    block: Block,
}

impl Parse for ArgGenerator {
    fn parse(input: ParseStream) -> syn::parse::Result<Self> {
        Ok(ArgGenerator {
            left_or: input.parse()?,
            pattern: input.parse()?,
            right_or: input.parse()?,
            block: Block {
                brace_token: Default::default(),
                stmts: Block::parse_within(input)?,
            },
        })
    }
}

#[derive(Default)]
struct YieldVisitor {
    found_exprs: usize,
    found_statement_exprs: usize,
    errors: Vec<(String, Span)>,
}

impl<'a> Visit<'a> for YieldVisitor {
    fn visit_stmt(&mut self, i: &'a Stmt) {
        if let Stmt::Semi(expr, _) = i {
            if let Expr::Yield(_) = expr {
                self.found_statement_exprs += 1
            }
        }

        visit::visit_stmt(self, i)
    }

    fn visit_expr_yield(&mut self, i: &'a ExprYield) {
        self.found_exprs += 1;

        visit::visit_expr_yield(self, i)
    }

    fn visit_expr_await(&mut self, i: &'a ExprAwait) {
        self.errors.push((
            format!("Await must not be used inside of a generator"),
            i.await_token.span,
        ))
    }

    fn visit_expr_async(&mut self, _i: &'a ExprAsync) {
        // Don't defer to the standard implementation.
        // This is so that `await`s within `async`s aren't
        // caught above.
    }

    fn visit_item_fn(&mut self, _i: &'a ItemFn) {
        // Likewise, we don't care about the contents
        // of locally defined functions.
    }
}

struct BlockVisitor {}

impl VisitMut for BlockVisitor {
    fn visit_expr_mut(&mut self, i: &mut Expr) {
        if let Expr::Yield(expr) = i {
            let yield_expr = expr
                .expr
                .take()
                .unwrap_or_else(|| Box::new(parse_quote! {()}));

            *i = Expr::Await(ExprAwait {
                attrs: std::mem::replace(&mut expr.attrs, vec![]),
                await_token: <Token!(await)>::default(),
                dot_token: <Token!(.)>::default(),
                base: parse_quote! {
                    __yield_awaiter(#yield_expr)
                },
            })
        }

        visit_mut::visit_expr_mut(self, i)
    }
}
