extern crate proc_macro;

use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::{format_ident, quote, ToTokens};
use rustc_hash::FxHashMap;
use std::iter;
use syn::parse::{Parse, ParseStream};
use syn::punctuated::Punctuated;
use syn::visit_mut::VisitMut;
use syn::{
    File, ForeignItem, ForeignItemFn, Ident, Item, PatType, Path, Signature, Token, TypeGroup, bracketed, parse_macro_input
};

const CUDA_RS: &'static str = include_str! {"cuda.rs"};
const NVML_RS: &'static str = include_str! {"nvml.rs"};
const CUBLAS_RS: &'static str = include_str! {"cublas.rs"};
const CUBLASLT_RS: &'static str = include_str! {"cublaslt.rs"};
const CUBLASLT_INTERNAL_RS: &'static str = include_str! {"cublaslt_internal.rs"};
const CUFFT_RS: &'static str = include_str! {"cufft.rs"};
const CUSPARSE_RS: &'static str = include_str! {"cusparse.rs"};
const CUDNN8_RS: &'static str = include_str! {"cudnn8.rs"};
const CUDNN9_RS: &'static str = include_str! {"cudnn9.rs"};
const ROCBLAS_RS: &'static str = include_str! {"../../ext/rocblas-sys/src/lib.rs"};
const ROCSPARSE_RS: &'static str = include_str! {"../../ext/rocsparse-sys/src/lib.rs"};
const HIPBLASLT_RS: &'static str = include_str! {"../../ext/hipblaslt-sys/src/lib.rs"};
const MIOPEN_RS: &'static str = include_str! {"../../ext/miopen-sys/src/lib.rs"};

// This macro accepts following arguments:
// * `normal_macro`: ident for a normal macro
// * zero or more:
//   * `override_macro`: ident for an override macro
//   * `override_fns`: list of override functions
// Then macro goes through every function in rust.rs, and for every fn `foo`:
// * if `foo` is contained in `override_fns` then pass it into `override_macro`
// * if `foo` is not contained in `override_fns` pass it to `normal_macro`
// Both `override_macro` and `normal_macro` expect semicolon-separated list:
//   macro_foo!(
//      "system" fn cuCtxDetach(ctx: CUcontext) -> CUresult;
//      "system" fn cuCtxDetach(ctx: CUcontext) -> CUresult
//   )
// Additionally, it does a fixup of CUDA types so they get prefixed with `type_path`
#[proc_macro]
pub fn cuda_function_declarations(tokens: TokenStream) -> TokenStream {
    function_declarations(tokens, CUDA_RS, false)
}

#[proc_macro]
pub fn cublas_function_declarations(tokens: TokenStream) -> TokenStream {
    function_declarations(tokens, CUBLAS_RS, false)
}

#[proc_macro]
pub fn cublaslt_function_declarations(tokens: TokenStream) -> TokenStream {
    function_declarations(tokens, CUBLASLT_RS, false)
}

#[proc_macro]
pub fn cublaslt_internal_function_declarations(tokens: TokenStream) -> TokenStream {
    function_declarations(tokens, CUBLASLT_INTERNAL_RS, false)
}

#[proc_macro]
pub fn cufft_function_declarations(tokens: TokenStream) -> TokenStream {
    function_declarations(tokens, CUFFT_RS, false)
}

#[proc_macro]
pub fn cusparse_function_declarations(tokens: TokenStream) -> TokenStream {
    function_declarations(tokens, CUSPARSE_RS, false)
}

#[proc_macro]
pub fn cudnn8_function_declarations(tokens: TokenStream) -> TokenStream {
    function_declarations(tokens, CUDNN8_RS, false)
}

#[proc_macro]
pub fn cudnn9_function_declarations(tokens: TokenStream) -> TokenStream {
    function_declarations(tokens, CUDNN9_RS, false)
}

#[proc_macro]
pub fn rocblas_function_declarations(tokens: TokenStream) -> TokenStream {
    function_declarations(tokens, ROCBLAS_RS, true)
}

#[proc_macro]
pub fn rocsparse_function_declarations(tokens: TokenStream) -> TokenStream {
    function_declarations(tokens, ROCSPARSE_RS, true)
}

#[proc_macro]
pub fn hipblaslt_function_declarations(tokens: TokenStream) -> TokenStream {
    function_declarations(tokens, HIPBLASLT_RS, true)
}

#[proc_macro]
pub fn miopen_function_declarations(tokens: TokenStream) -> TokenStream {
    function_declarations(tokens, MIOPEN_RS, true)
}

fn function_declarations(tokens: TokenStream, module: &str, relaxed: bool) -> TokenStream {
    let input = parse_macro_input!(tokens as FnDeclInput);
    let mut cuda_module = syn::parse_str::<File>(module).unwrap();
    let mut choose_macro = ChooseMacro::new(input);
    syn::visit_mut::visit_file_mut(&mut FixFnSignatures, &mut cuda_module);
    for item in cuda_module.items {
        let extern_ = if let Item::ForeignMod(extern_) = item {
            extern_
        } else if relaxed {
            continue;
        } else {
            unreachable!()
        };
        let abi = extern_.abi.name;
        for mut item in extern_.items {
            if let ForeignItem::Fn(ForeignItemFn {
                sig: Signature { ref ident, ref mut output, ref variadic, .. },
                ref mut attrs,
                ref mut vis,
                ..
            }) = item
            {
                if relaxed {
                    if variadic.is_some() {
                        continue;
                    }
                    *vis = syn::Visibility::Inherited;
                    if output == &syn::ReturnType::Default {
                        *output = syn::ReturnType::Type(
                            Token![->](Span::call_site()),
                            Box::new(syn::parse_quote! { () }),
                        );
                    }
                }
                *attrs = Vec::new();
                choose_macro.add(ident, quote! { #abi #item });
            } else {
                unreachable!()
            }
        }
    }
    let mut result = proc_macro2::TokenStream::new();
    for (path, items) in
        iter::once(choose_macro.default).chain(choose_macro.override_sets.into_iter())
    {
        if items.is_empty() {
            continue;
        }
        quote! {
            #path ! { #(#items)* }
        }
        .to_tokens(&mut result);
    }
    result.into()
}

#[proc_macro]
pub fn nvml_function_declarations(tokens: TokenStream) -> TokenStream {
    function_declarations(tokens, NVML_RS, false)
}
struct FnDeclInput {
    normal_macro: Path,
    overrides: Punctuated<OverrideMacro, Token![,]>,
}

impl Parse for FnDeclInput {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let normal_macro = input.parse::<Path>()?;
        let overrides = if input.is_empty() {
            Punctuated::new()
        } else {
            input.parse::<Token![,]>()?;
            input.parse_terminated(OverrideMacro::parse, Token![,])?
        };
        Ok(Self {
            normal_macro,
            overrides,
        })
    }
}
struct OverrideMacro {
    macro_: Path,
    functions: Punctuated<Ident, Token![,]>,
}

impl Parse for OverrideMacro {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let macro_ = input.parse::<Path>()?;
        input.parse::<Token![<=]>()?;
        let functions_content;
        bracketed!(functions_content in input);
        let functions = functions_content.parse_terminated(Ident::parse, Token![,])?;
        Ok(Self { macro_, functions })
    }
}

struct ChooseMacro {
    default: (Path, Vec<proc_macro2::TokenStream>),
    override_lookup: FxHashMap<Ident, Path>,
    override_sets: FxHashMap<Path, Vec<proc_macro2::TokenStream>>,
}

impl ChooseMacro {
    fn new(input: FnDeclInput) -> Self {
        let mut override_lookup = FxHashMap::default();
        let mut override_sets = FxHashMap::default();
        for OverrideMacro { macro_, functions } in input.overrides {
            for ident in functions {
                override_lookup.insert(ident, macro_.clone());
                override_sets.insert(macro_.clone(), Vec::new());
            }
        }
        Self {
            default: (input.normal_macro, Vec::new()),
            override_lookup,
            override_sets,
        }
    }

    fn add(&mut self, ident: &Ident, tokens: proc_macro2::TokenStream) {
        match self.override_lookup.get(ident) {
            Some(override_macro) => {
                self.override_sets
                    .get_mut(override_macro)
                    .unwrap()
                    .push(tokens);
            }
            None => self.default.1.push(tokens),
        }
    }
}

// For some reason prettyplease will append trailing comma *only*
// if there are two or more arguments
struct FixFnSignatures;

impl VisitMut for FixFnSignatures {
    fn visit_signature_mut(&mut self, s: &mut syn::Signature) {
        s.inputs.pop_punct();
    }
}

const MODULES: &[&str] = &[
    "array", "context", "device", "driver", "event", "function", "graph", "kernel",
    "library", "link", "memory", "module", "pointer", "stream", "tex"
];

fn normalize_fn_impl(
    prefix: &str,
    default_module: Option<&str>,
    tokens: TokenStream,
) -> TokenStream {
    let mut path = parse_macro_input!(tokens as syn::Path);
    let fn_ = path
        .segments
        .pop()
        .unwrap()
        .into_tuple()
        .0
        .ident
        .to_string();
    let already_has_module = MODULES.contains(&&*path.segments.last().unwrap().ident.to_string());
    let segments: Vec<String> = split(&fn_[prefix.len()..]); // skip "cu"
    let fn_path = join(segments, default_module.filter(|_| !already_has_module));
    quote! {
        #path #fn_path
    }
    .into()
}

#[proc_macro]
pub fn cuda_normalize_fn(tokens: TokenStream) -> TokenStream {
    normalize_fn_impl("cu", Some("driver"), tokens)
}

#[proc_macro]
pub fn cublas_normalize_fn(tokens: TokenStream) -> TokenStream {
    normalize_fn_impl("cublas", None, tokens)
}

#[proc_macro]
pub fn cublaslt_normalize_fn(tokens: TokenStream) -> TokenStream {
    normalize_fn_impl("cublasLt", None, tokens)
}

#[proc_macro]
pub fn cudnn_normalize_fn(tokens: TokenStream) -> TokenStream {
    normalize_fn_impl("cudnn", None, tokens)
}

#[proc_macro]
pub fn cusparse_normalize_fn(tokens: TokenStream) -> TokenStream {
    normalize_fn_impl("cusparse", None, tokens)
}

#[proc_macro]
pub fn nvml_normalize_fn(tokens: TokenStream) -> TokenStream {
    normalize_fn_impl("nvml", None, tokens)
}

#[proc_macro]
pub fn generate_input_struct(tokens: TokenStream) -> TokenStream {
    let input = parse_macro_input!(tokens as FnWithFields);
    let struct_name = format_ident!("{}In", input.fn_name);
    let fields = input.fields.iter().filter(|pat| is_input_field(&pat.ty)).map(use_wire_object);
    quote! {
        #[derive(Archive, Deserialize, Serialize, Debug, PartialEq)]
        #[repr(C)]
        pub struct #struct_name {
            #(#fields),*
        }
        unsafe impl Portable for #struct_name {}
    }.into()
}

#[proc_macro]
pub fn generate_output_struct(tokens: TokenStream) -> TokenStream {
    let input = parse_macro_input!(tokens as FnWithFields);
    let struct_name = format_ident!("{}Out", input.fn_name);
    let fields = input.fields.iter().filter(|pat| !is_input_field(&pat.ty)).map(strip_mut_ptr);
    quote! {
        #[derive(Archive, Deserialize, Serialize, Debug, PartialEq, Clone)]
        #[repr(C)]
        pub struct #struct_name {
            #(#fields),*
        }
        unsafe impl Portable for #struct_name {}
    }.into()
}

fn is_input_field(ty: &syn::Type) -> bool {
    match ty {
        syn::Type::Ptr(token) if token.mutability.is_some() => false,
        syn::Type::Group(TypeGroup {elem, ..}) => is_input_field(&*elem),
        _ => true
    }
}

fn use_wire_object(ty: &syn::PatType) -> proc_macro2::TokenStream {
    let pat = &ty.pat;
    let type_ = &ty.ty;
    quote! { pub #pat : <#type_ as crate::CudaEncode>::WireObject }
}

fn strip_mut_ptr(ty: &syn::PatType) -> proc_macro2::TokenStream {
    fn strip_mut_ptr_from_type(ty: &syn::Type) -> proc_macro2::TokenStream {
        match ty {
            syn::Type::Ptr(token) => {
                let elem = &token.elem;
                quote! { #elem }
            },
            _ => unreachable!()
        }
    }
    let pat = &ty.pat;
    match &*ty.ty {
        syn::Type::Ptr(token) => {
            let type_ = strip_mut_ptr_from_type(&*token.elem);
            quote! { pub #pat : <#type_ as crate::CudaEncode>::WireObject }
        },
        syn::Type::Group(TypeGroup {elem, ..}) => {
            let type_ = strip_mut_ptr_from_type(&*elem);
            quote! { pub #pat : <#type_ as crate::CudaEncode>::WireObject }
        },
        _ => unreachable!()
    }
}

struct FnWithFields {
    fn_name: Ident,
    fields: Punctuated<PatType, Token![,]>,
}

impl Parse for FnWithFields {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let fn_name = input.parse::<Ident>()?;
        let fields = if input.is_empty() {
            Punctuated::new()
        } else {
            input.parse::<Token![,]>()?;
            input.parse_terminated(PatType::parse, Token![,])?
        };
        Ok(Self {
            fn_name,
            fields,
        })
    }
}

fn split(fn_: &str) -> Vec<String> {
    let mut result = Vec::new();
    let mut chars: iter::Peekable<_> = fn_.chars().peekable();
    while let Some(c) = chars.next() {
        // Special handling of 2D/3D etc.
        if c.is_ascii_digit() && chars.peek() == Some(&'D') {
            chars.next();
            result.push(format!("{c}d"));
            continue;
        }
        if c.is_ascii_uppercase() {
            result.push(c.to_ascii_lowercase().to_string());
        } else {
            result.last_mut().unwrap().push(c);
        }
    }
    result
}

fn join(
    fn_: Vec<String>,
    default_module: Option<&str>,
) -> Punctuated<Ident, Token![::]> {
    fn full_form(segment: &str) -> Option<&[&str]> {
        Some(match segment {
            "ctx" => &["context"],
            "func" => &["function"],
            "mem" => &["memory"],
            "memcpy" => &["memory", "copy"],
            "memset" => &["memory", "set"],
            _ => return None,
        })
    }
    fn safe_ident(segment: &str) -> Ident {
        if segment.chars().next().unwrap().is_ascii_digit() {
            Ident::new(&format!("_{segment}"), Span::call_site())
        } else {
            Ident::new(segment, Span::call_site())
        }
    }
    let mut normalized: Vec<&str> = Vec::new();
    for segment in fn_.iter() {
        match full_form(segment) {
            Some(segments) => normalized.extend(segments.into_iter()),
            None => normalized.push(&*segment),
        }
    }
    if let Some(default_module) = default_module {
        if !MODULES.contains(&normalized[0]) {
            let mut globalized = vec![default_module];
            globalized.extend(normalized);
            normalized = globalized;
        }
        let (module, path) = normalized.split_first().unwrap();
        let path = path.join("_");
        [module, &&*path]
            .into_iter()
            .map(|s| safe_ident(s))
            .collect()
    } else {
        [safe_ident(&normalized.join("_"))]
            .into_iter()
            .collect()
    }

}

#[proc_macro_attribute]
pub fn test_cuda(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let fn_ = parse_macro_input!(item as syn::ItemFn);
    let cuda_fn = format_ident!("{}{}", fn_.sig.ident, "_nvidia");
    let zluda_fn = format_ident!("{}{}", fn_.sig.ident, "_zluda");
    let fn_name = fn_.sig.ident.clone();
    quote! {
        #[test]
        fn #cuda_fn() {
            unsafe { #fn_name(<crate::tests::Cuda>::new()) }
        }
        #[test]
        fn #zluda_fn() {
            unsafe { #fn_name(<crate::tests::Zluda>::new()) }
        }

        #fn_
    }.into()
}
