// pest-typed. A statically typed version of pest.
// Copyright (c) 2023 黄博奕
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use proc_macro2::{Ident, TokenStream};
use quote::{format_ident, quote, ToTokens};
use syn::{
    braced,
    parse::{Parse, ParseStream},
    Expr, Token,
};

struct MatchChoice {
    ident: Ident,
    body: Expr,
}
impl Parse for MatchChoice {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let ident: Ident = input.parse()?;
        let _: Token!(=>) = input.parse()?;
        let body: Expr = input.parse()?;
        let _: Option<Token!(,)> = input.parse()?;
        Ok(Self { ident, body })
    }
}
impl MatchChoice {
    fn to_tokens(&self, tag: Option<&TokenStream>) -> TokenStream {
        let ident = &self.ident;
        let body = &self.body;
        if let Some(tag) = tag {
            quote! {
                #tag(#ident) => #body,
            }
        } else {
            quote! {
                #ident => #body
            }
        }
    }
}

struct MatchChoices {
    expr: Box<Expr>,
    arms: Vec<MatchChoice>,
}
impl Parse for MatchChoices {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let expr = Expr::parse_without_eager_brace(input)?;

        let content;
        let _brace_token = braced!(content in input);

        let mut arms = Vec::new();
        while !content.is_empty() {
            arms.push(content.parse()?);
        }

        Ok(Self {
            expr: Box::new(expr),
            arms,
        })
    }
}
impl ToTokens for MatchChoices {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let expr = &self.expr;
        let n = self.arms.len();
        let t = format_ident!("Choice{}", n);
        let t = quote! {generics::#t};
        let arms = &self
            .arms
            .iter()
            .enumerate()
            .map(|(i, arm)| {
                if n > 1 {
                    let variant = format_ident!("_{}", i);
                    let t = quote!(#t::#variant);
                    arm.to_tokens(Some(&t))
                } else {
                    arm.to_tokens(None)
                }
            })
            .collect::<Vec<_>>();
        tokens.extend(
            quote! {
                match #expr {
                    #(
                        #arms
                    )*
                }
            }
            .into_iter(),
        );
    }
}

/// Match a choices structure generated by pest-typed.
///
/// Please imports generated module `generics` so that the macro can use it.
pub fn match_choices(input: TokenStream) -> TokenStream {
    let mc: MatchChoices = syn::parse2(input).unwrap();
    quote! {#mc}.into()
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn choice_0() {
        let input = quote! {
            x {}
        }
        .into();
        let _: TokenStream = match_choices(input);
    }
    #[test]
    fn choice_1() {
        let input = quote! {
            x {
                a => {
                    println!("{a}");
                }
            }
        }
        .into();
        let _: TokenStream = match_choices(input);
    }
    #[test]
    fn choice_2() {
        let input = quote! {
            x {
                a => a - 1,
                b => b + 1
            }
        }
        .into();
        let _: TokenStream = match_choices(input);
    }
}
