use proc_macro2::{Ident, TokenStream};
use quote::{format_ident, quote, ToTokens, TokenStreamExt};
use syn::FnArg;

use super::c3_ast::{ClassDef, ClassFnImpl, ClassNameDef, FnDef, PackageDef, VarDef};

impl ToTokens for PackageDef {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        tokens.append_all(&self.other_code);
        tokens.extend(stack_definition());
        tokens.extend(self.class_name.to_token_stream());
        tokens.append_all(&self.classes);
    }
}

impl ToTokens for ClassNameDef {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        let classes = &self.classes;
        tokens.extend(quote! {
            #[derive(Clone)]
            enum ClassName {
                #(#classes),*
            }
        })
    }
}

impl ToTokens for ClassDef {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        let class_ident = &self.class;
        let path = &self.path;
        let path_len = path.len();
        let variables = &self.variables;
        let functions = &self.functions;
        tokens.extend(quote! {
            pub struct #class_ident {
                __stack: PathStack,
                #(#variables),*
            }

            impl #class_ident {
                const PATH: &'static [ClassName; #path_len] = &[
                    #(ClassName::#path),*
                ];

                #(#functions)*
            }
        })
    }
}

impl ToTokens for VarDef {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        let ident = &self.ident;
        let ty = &self.ty;
        tokens.extend(quote! {
            #ident: #ty
        });
    }
}

impl ToTokens for FnDef {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        let fn_ident = &self.name;
        let fn_super_ident = format_ident!("super_{}", fn_ident.to_string());
        let args = &self.args;
        let ret = &self.ret;
        let implementations = &self.implementations;
        let args_as_params = args_to_params(&args);
        tokens.extend(quote! {
            pub fn #fn_ident(#(#args),*) #ret {
                self.__stack.push_path_on_stack(Self::PATH);
                let result = self.#fn_super_ident(#(#args_as_params),*);
                self.__stack.drop_one_from_stack();
                result
            }

            pub fn #fn_super_ident(#(#args),*) #ret {
                let __class = self.__stack.pop_from_top_path();
                match __class {
                    #(#implementations),*
                    #[allow(unreachable_patterns)]
                    _ => self.#fn_super_ident(#(#args_as_params),*),
                }
            }
        });
    }
}

impl ToTokens for ClassFnImpl {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        let class = &self.class;
        let implementation = &self.implementation;
        tokens.extend(quote! {
            ClassName::#class => #implementation
        });
    }
}

fn args_to_params(args: &Vec<FnArg>) -> Vec<Ident> {
    args.clone()
        .into_iter()
        .skip(1)
        .map(|arg| {
            if let FnArg::Typed(arg) = arg {
                format_ident!("{}", arg.pat.to_token_stream().to_string())
            } else {
                panic!("Unsupported arg");
            }
        })
        .collect()
}

fn stack_definition() -> TokenStream {
    quote! {
        struct PathStack {
            stack: std::sync::Arc<std::sync::Mutex<Vec<Vec<ClassName>>>>
        }

        impl PathStack {
            pub fn new() -> Self {
                PathStack {
                    stack: std::sync::Arc::new(std::sync::Mutex::new(Vec::new()))
                }
            }

            pub fn push_path_on_stack(&self, path: &[ClassName]) {
                let mut stack = self.stack.lock().unwrap();
                stack.push(path.to_vec());
            }

            pub fn drop_one_from_stack(&self) {
                let mut stack = self.stack.lock().unwrap();
                stack.pop().unwrap();
            }

            pub fn pop_from_top_path(&self) -> ClassName {
                let mut stack = self.stack.lock().unwrap();
                let mut path = stack.pop().unwrap();
                let class = path.pop().unwrap();
                stack.push(path);
                class
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use quote::{quote, ToTokens};

    use crate::{c3_ast_builder::tests::test_c3_ast, test_utils::test_code};

    use super::stack_definition;

    #[test]
    fn test_package_def_printing() {
        let input = test_c3_ast();
        let stack = stack_definition();
        let target = quote! {
            pub type Num = u32;

            #stack

            #[derive(Clone)]
            enum ClassName {
                A,
                B,
            }
            pub struct B {
                __stack: PathStack,
                x: u32,
            }
            impl B {
                const PATH: &'static [ClassName; 2usize] = &[ClassName::B, ClassName::A];
                pub fn bar(&self, counter: Num) -> String {
                    self.__stack.push_path_on_stack(Self::PATH);
                    let result = self.super_bar(counter);
                    self.__stack.drop_one_from_stack();
                    result
                }
                pub fn super_bar(&self, counter: Num) -> String {
                    let __class = self.__stack.pop_from_top_path();
                    match __class {
                        ClassName::A => {
                            let label = format!("A::bar({})", counter);
                            if counter == 0 {
                                label
                            } else {
                                format!("{} {}", label, self.foo(counter - 1))
                            }
                        }
                        ClassName::B => {
                            let label = format!("B::bar({})", counter);
                            if counter == 0 {
                                label
                            } else {
                                format!("{} {}", label, self.super_bar(counter - 1))
                            }
                        }
                        #[allow(unreachable_patterns)]
                        _ => self.super_bar(counter),
                    }
                }
                pub fn foo(&self, counter: Num) -> String {
                    self.__stack.push_path_on_stack(Self::PATH);
                    let result = self.super_foo(counter);
                    self.__stack.drop_one_from_stack();
                    result
                }
                pub fn super_foo(&self, counter: Num) -> String {
                    let __class = self.__stack.pop_from_top_path();
                    match __class {
                        ClassName::A => {
                            let label = format!("A::foo({})", counter);
                            if counter == 0 {
                                label
                            } else {
                                format!("{} {}", label, self.bar(counter - 1))
                            }
                        }
                        #[allow(unreachable_patterns)]
                        _ => self.super_foo(counter),
                    }
                }
            }

        };
        test_code(input.to_token_stream(), target)
    }
}