#![feature(box_patterns)]

//! DO NOT USE THIS CRATE DIRECTLY
//! 
//! provides procedural macros for rs2glsl

use proc_macro::TokenStream;
use proc_macro2::{Ident, TokenStream as TokenStream2};
use quote::{quote, quote_spanned, TokenStreamExt};
use std::collections::HashSet;
use syn::{
    parse_macro_input,
    spanned::Spanned,
    visit::{self, Visit},
    ExprPath, FnArg, Item, ItemConst, ItemFn, ItemMod, Local, Pat,
};

mod generator;
mod pretty;
mod translator;
use translator::*;

#[derive(Default, Debug)]
struct DependencyCollector {
    /// definitions in local scope
    local_defs: HashSet<String>,
    global_refs: HashSet<Ident>,
}

impl<'ast> Visit<'ast> for DependencyCollector {
    fn visit_fn_arg(&mut self, i: &'ast FnArg) {
        match i {
            FnArg::Receiver(_) => self.local_defs.insert("self".into()),
            FnArg::Typed(t) => match &*t.pat {
                Pat::Ident(i) => self.local_defs.insert(i.ident.to_string()),
                _ => true,
            },
        };
        visit::visit_fn_arg(self, i)
    }
    fn visit_local(&mut self, i: &'ast Local) {
        match &i.pat {
            Pat::Type(i) => match &i.pat {
                box Pat::Ident(i) => {
                    self.local_defs.insert(i.ident.to_string());
                }
                _ => (),
            },
            Pat::Ident(i) => {
                self.local_defs.insert(i.ident.to_string());
            }
            _ => (),
        };
        visit::visit_local(self, i)
    }
    fn visit_expr_path(&mut self, i: &'ast ExprPath) {
        if let Some(i) = i.path.get_ident() {
            if !self.local_defs.contains(&i.to_string()) {
                self.global_refs.insert(i.clone());
            }
        }
        visit::visit_expr_path(self, i)
    }
}

impl DependencyCollector {
    fn collect_dependencies(fun: &ItemFn) -> Vec<Ident> {
        let mut collector = Self::default();
        collector.visit_item_fn(fun);
        let mut refs = collector
            .global_refs
            .into_iter()
            .filter(|id| !generator::is_builtin(&id.to_string()))
            .collect::<Vec<_>>();
        refs.dedup_by_key(|i| i.to_string());
        refs
    }
}

macro_rules! get_vulkan_def {
    ($item:expr) => {{
        let mut vulkan_def = String::new();
        match $item.to_vulkan(&mut vulkan_def) {
            Err(ToVulkanError::Unsupported(span, reason)) => {
                println!("{vulkan_def}");
                let mut error_msg = String::from("vulkan_fn: unsupported feature");
                match reason {
                    Reason::Unknown => {}
                    other_reason => {
                        error_msg += " (";
                        error_msg += other_reason.as_str();
                        error_msg += ")";
                    }
                }
                return quote_spanned! { span =>
                    compile_error!(#error_msg);
                }
                .into();
            }
            Err(ToVulkanError::FmtErr(err)) => panic!("unexpected std::fmt error: {err:?}"),
            Ok(()) => (),
        };
        let vulkan_def = pretty::reformat(&vulkan_def);
        vulkan_def
    }};
}

fn glsl_fn(tokens: TokenStream, omit_impl: bool) -> TokenStream {
    let fun = parse_macro_input!(tokens as ItemFn);

    let fn_name = fun.sig.ident.clone();
    let name_as_str = fn_name.to_string();
    let vis = fun.vis.clone();
    let struct_name = struct_name(&fn_name);
    let glsl_item = item_name(&fn_name);

    let fn_dependencies;
    let vulkan_def;

    if omit_impl {
        fn_dependencies = quote! {};
        vulkan_def = String::new();
    } else {
        fn_dependencies = DependencyCollector::collect_dependencies(&fun)
            .into_iter()
            .map(|i| {
                let i = item_name(&i);
                quote! { Box::new(#i), }
            })
            .collect::<TokenStream2>();

        vulkan_def = get_vulkan_def!(fun);
    }

    quote! {
        // original function
        #fun

        // type that stores glsl metadata
        #[allow(non_camel_case_types)]
        #[doc(hidden)]
        #[derive(Copy, Clone, Debug)]
        #vis struct #struct_name;
        impl ::rs2glsl::GlslItem for #struct_name {
            fn name(&self) -> &'static str {
                #name_as_str
            }
            fn dependencies(&self) -> Vec<Box<dyn ::rs2glsl::GlslItem>> {
                vec![#fn_dependencies]
            }
            fn definition(&self) -> &'static str {
                #vulkan_def
            }
        }
        #vis const #glsl_item: #struct_name = #struct_name;
    }
    .into()
}

fn glsl_const(tokens: TokenStream, omit_impl: bool) -> TokenStream {
    let input = parse_macro_input!(tokens as ItemConst);
    let const_name = input.ident.clone();
    let struct_name = struct_name(&const_name);
    let vis = input.vis.clone();
    let name_as_str = const_name.to_string();
    let glsl_item = item_name(&const_name);

    let vulkan_def = if omit_impl {
        String::new()
    } else {
        get_vulkan_def!(input)
    };

    quote! {
        // original constant
        #input

        // glsl item
        #[allow(non_camel_case_types)]
        #[doc(hidden)]
        #[derive(Copy, Clone, Debug)]
        #vis struct #struct_name;
        impl ::rs2glsl::GlslItem for #struct_name {
            fn name(&self) -> &'static str {
                #name_as_str
            }
            fn dependencies(&self) -> Vec<Box<dyn ::rs2glsl::GlslItem>> {
                vec![]
            }
            fn definition(&self) -> &'static str {
                #vulkan_def
            }
        }

        #vis const #glsl_item: #struct_name = #struct_name;
    }
    .into()
}

fn glsl_mod(tokens: TokenStream, omit_impl: bool) -> TokenStream {
    let input = parse_macro_input!(tokens as ItemMod);
    let input_span = input.span();
    let mod_name = input.ident;

    let result = match input.content {
        Some((_, items)) => {
            let mut tokens = TokenStream2::new();
            for item in items.iter() {
                let item = quote! {#item}.into();
                let item: TokenStream2 = inner_glsl(item, omit_impl).into();
                tokens.append_all(item);
            }
            quote! {mod #mod_name{#tokens}}
        }
        None => quote! {compile_error!("Cannot annotate modules of other files.")},
    };
    quote_spanned! {input_span => #result}.into()
}

fn inner_glsl(tokens: TokenStream, omit_impl: bool) -> TokenStream {
    let tokens2 = tokens.clone();
    let item = parse_macro_input!(tokens2 as Item);
    let handler = match item {
        Item::Const(_) => glsl_const,
        Item::Fn(_) => glsl_fn,
        Item::Mod(_) => glsl_mod,
        _ => return tokens,
    };
    handler(tokens, omit_impl)
}

/// Annotates an item in Rust where a corresponding item in GLSL should be created
/// 
/// Only a subset of the usual Rust syntax wil. work.
#[proc_macro_attribute]
pub fn glsl(_args: TokenStream, tokens: TokenStream) -> TokenStream {
    inner_glsl(tokens, false)
}

/// Define a function or constant that is present in the GLSL context.
/// 
/// This will not create a GLSL implementation but assumes it already exists.
/// Full Rust syntax can be used.
#[proc_macro_attribute]
pub fn glsl_native(_args: TokenStream, tokens: TokenStream) -> TokenStream {
    inner_glsl(tokens, true)
}

/// Generates the GLSL types in Rust
#[proc_macro]
pub fn make_types(_: TokenStream) -> TokenStream {
    let implementation = generator::generate();
    implementation.into()
}
