extern crate proc_macro;
extern crate syn;
extern crate quote;
extern crate proc_macro2;

use proc_macro::TokenStream;
use quote::quote;
use proc_macro2::Span;
use syn::{braced, bracketed, parenthesized, parse::{Parse, ParseStream}, parse_macro_input, punctuated::Punctuated, token::{self, Paren}, Block, Error, PatType, Ident, LitInt, Result, Token, Type, TypeTuple};

#[proc_macro]
pub fn generate_opcodes(item: TokenStream) -> TokenStream {
    let ops_def = parse_macro_input!(item as OpcodeBlock);

    let enum_name = ops_def.enum_name;
    let field_names = ops_def.codes.iter().map(|field| &field.name).collect::<Vec<_>>();

    let getcode_match_arms = ops_def.codes.iter().map(|field| {
        let code = &field.byte_code;
        let variant = &field.name;
        quote! {
            #code => Some(Self::#variant),
        }
    }).collect::<Vec<_>>();

    let getinput_match_arms = ops_def.codes.iter().map(|field| {
        let variant = &field.name;
        let num = &field.input_params.len();
        quote! {
            #enum_name::#variant => #num,
        }
    });

    let getoutput_match_arms = ops_def.codes.iter().map(|field| {
        let variant = &field.name;
        let num = &field.output_params.len();
        quote! {
            #enum_name::#variant => #num,
        }
    });

    let gen_stub_arms = ops_def.codes.iter().map(|field| {
        let variant = &field.name;
        let value = &field.no_store;
        quote! {
            #enum_name::#variant => #value,
        }
    });

    let closure_match_arms = ops_def.codes.iter().map(|field| {
        let variant = &field.name;
        let closure_expr = &field.func_block;
        let interpreter_arg = &field.interpreter_ident.clone().unwrap_or(Ident::new("_terp", Span::call_site()));
        let input_args = &field.input_params;
        let output_args = if field.output_params.len() == 1 {
            let arg = field.output_params[0].clone();
            quote! { #arg }
        } else {
            let args = &field.output_params;
            quote! { (#(#args),*) }
        };


        let call_args = if field.input_params.len() > 0 {
            let arg_expr = (0..field.input_params.len()).map(|index| quote! {input_params[#index]}).collect::<Vec<_>>();
            quote! {interpreter, #(#arg_expr),*}
        } else {
            quote! {interpreter}
        };

        let return_vec = (0..field.output_params.len()).map(|index| quote! {out_tuple.#index});
        let call = match field.output_params.len() {
            0 => quote! {
                closure(#call_args);
                return Vec::new();
            },
            1 => quote! {
                let result = closure(#call_args);
                return vec![ result ];
            },
            _ => quote! {
                let out_tuple = closure(#call_args);
                return vec![ #(#return_vec),* ];
            },
        };

        quote! {
            #enum_name::#variant => {
                let closure = |#interpreter_arg: &mut crate::Interpreter, #(#input_args),*| -> #output_args
                    #closure_expr;

                #call
            }
        }
    });

    let tokens = quote! {
        #[derive(Debug, Clone)]
        pub enum #enum_name {
            #(#field_names),*
        }

        impl #enum_name {
            pub fn get_from_code(code: u32) -> Option<Self> {
                match code {
                    #(#getcode_match_arms)*
                    _ => None
                }
            }

            pub fn get_input_arg_count(opcode: Self) -> usize {
                match opcode {
                    #(#getinput_match_arms)*
                }
            }

            pub fn get_output_arg_count(opcode: Self) -> usize {
                match opcode {
                    #(#getoutput_match_arms)*
                }
            }

            pub fn should_gen_return_stub(opcode: Self) -> bool {
                match opcode {
                    #(#gen_stub_arms)*
                }
            }

            pub fn call_opcode(interpreter: &mut crate::Interpreter, opcode: Self, input_params: Vec<u32>) -> Vec<u32> {
                match opcode {
                    #(#closure_match_arms)*
                }
            }
        }
    };

    tokens.into()
}

impl Parse for OpcodeBlock {
    fn parse(input: ParseStream) -> Result<Self> {
        let inner;
        input.parse::<Token![pub]>()?;
        Ok(OpcodeBlock {
            enum_token: input.parse()?,
            enum_name: input.parse()?,
            brace_token: braced!(inner in input),
            codes: inner.parse_terminated(SingleOpcodeDefinition::parse, Token![,])?,
        })
    }
}

impl Parse for SingleOpcodeDefinition {
    fn parse(input: ParseStream) -> Result<Self> {
        let attr;
        input.parse::<Token![#]>()?;
        bracketed!(attr in input);
        let code = attr.parse::<Ident>()?;
        if code != "code" {
            return Err(Error::new(code.span(), "expected `code` attribute"));
        }
        let code;
        parenthesized!(code in attr);
        let byte_code = code.parse::<LitInt>()?;

        let mut no_store = false;
        if attr.peek(Token![,]) {
            attr.parse::<Token![,]>()?;
            let extra_attr = attr.parse::<Ident>()?;
            if extra_attr == "no_store" {
                no_store = true;
            }
        }

        input.parse::<Token![fn]>()?;
        let name = input.parse::<Ident>()?;

        let params;
        let mut interpreter_ident = None;
        parenthesized!(params in input);
        if params.peek(Token![&]) {
            params.parse::<Token![&]>()?;
            params.parse::<Token![mut]>()?;
            interpreter_ident = Some(params.parse::<Ident>()?);

            if params.peek(Token![,]) {
                params.parse::<Token![,]>()?;
            }
        }

        let input_params = params.parse_terminated(PatType::parse, Token![,])?
            .into_iter().map(|param| param).collect();

        let output_params;
        if input.peek(Token![->]) {
            input.parse::<Token![->]>()?;

            if input.peek(Paren) {
                let params;
                parenthesized!(params in input);
                output_params = params.parse_terminated(Type::parse, Token![,])?.into_iter().map(|s| s).collect::<Vec<_>>();
            } else {
                output_params = vec![input.parse::<Type>()?]
            }
        } else {
            output_params = Vec::new();
        }

        let func_block = input.parse::<Block>()?;

        Ok(SingleOpcodeDefinition {
            byte_code,
            no_store,
            name,
            input_params,
            interpreter_ident,
            output_params,
            func_block,
        })
    }
}

struct OpcodeBlock {
    pub enum_token: token::Enum,
    pub enum_name: Ident,
    pub brace_token: token::Brace,
    pub codes: Punctuated<SingleOpcodeDefinition, Token![,]>,
}

struct SingleOpcodeDefinition {
    pub byte_code: LitInt,
    pub no_store: bool,
    pub name: Ident,
    pub interpreter_ident: Option<Ident>,
    pub input_params: Vec<PatType>,
    pub output_params: Vec<Type>,
    pub func_block: Block,
}