use proc_macro2::{Ident, Span, TokenStream};
use quote::{quote, TokenStreamExt};
use cartesian::*;

pub(crate) fn generate() -> TokenStream {
    let vec = make_vectors();
    let mat = make_mats();
    quote! {
        #vec
        #mat
    }
}

pub fn is_builtin(item: &str) -> bool {
    for ty in TYS {
        if ty.defines(item) {
            return true;
        }
    }
    false
}

#[derive(Copy, Clone, Debug)]
struct Kind {
    /// prefix in the glsl types
    prefix: &'static str,

    /// rust type
    ty: &'static str,
}
impl Kind {
    const fn new(prefix: &'static str, ty: &'static str) -> Self {
        Self { prefix, ty }
    }
    fn has_arithmetic(self) -> bool {
        match self.ty {
            "bool" => false,
            _ => true,
        }
    }
    fn has_negation(self) -> bool {
        match self.ty {
            "bool" | "u32" => false,
            _ => true,
        }
    }

    /// name of the vector type with the specified dimension
    fn vec_name(self, dimensions: usize) -> String {
        format!("{}vec{dimensions}", self.prefix)
    }

    /// name of the square matrix type with the specified number of rows & cols
    fn sq_mat_name(self, dimensions: usize) -> String {
        format!("{}mat{dimensions}", self.prefix)
    }

    fn mat_name(self, cols: usize, rows: usize) -> String {
        format!("{}mat{cols}x{rows}", self.prefix)
    }

    fn member_ty_ident(self) -> Ident {
        Ident::new(self.ty, Span::call_site())
    }

    fn defines(self, item: &str) -> bool {
        for dim in DIMS {
            if self.vec_name(dim) == item || self.sq_mat_name(dim) == item {
                return true;
            }
        }
        for (&cols, &rows) in cartesian!(DIMS.iter(), DIMS.iter()) {
            if self.mat_name(cols, rows) == item {
                return true;
            }
        }
        false
    }
}

const TYS: [Kind; 5] = [
    Kind::new("b", "bool"),
    Kind::new("", "f32"),
    Kind::new("d", "f64"),
    Kind::new("i", "i32"),
    Kind::new("u", "u32"),
];
const DIMS: [usize; 3] = [2, 3, 4];

fn make_vectors() -> TokenStream {
    let mut impls = quote! {
        use core::ops::*;
        use crate::GlslItem;
    };

    let id = |ident| Ident::new(ident, Span::call_site());

    let axis = ["x", "y", "z", "w"];
    let cwise_binops = [("Add", "add"), ("Sub", "sub"), ("Mul", "mul")];
    let cwise_unops = [("Neg", "neg")];

    for len in DIMS {
        for ty in TYS {
            let ty_name_str = ty.vec_name(len);
            let ty_name = Ident::new(&ty_name_str, Span::call_site());
            let member_ty = ty.member_ty_ident();
            let coord = axis[..len].iter().map(|s| id(s)).collect::<Vec<_>>();

            impls.append_all(quote! {
                #[derive(Copy, Clone, Debug)]
                pub struct #ty_name {#(
                    pub #coord: #member_ty,
                )*}
                pub fn #ty_name(#(#coord: #member_ty,)*) -> #ty_name {
                    #ty_name {
                        #(#coord,)*
                    }
                }
            });

            if ty.has_arithmetic() {
                impls.append_all(quote! {
                    impl Mul<#member_ty> for #ty_name {
                        type Output = Self;
                        fn mul(self, rhs: #member_ty) -> Self {
                            Self {
                                #(#coord: self.#coord * rhs,)*
                            }
                        }
                    }
                });
                for (t, f) in cwise_binops {
                    let (t, f) = (id(t), id(f));
                    impls.append_all(quote! {
                        impl #t for #ty_name {
                            type Output = Self;
                            fn #f(self, other: Self) -> Self {
                                Self {
                                    #(#coord: self.#coord.#f(other.#coord),)*
                                }
                            }
                        }
                    });
                }
            }
            if ty.has_negation() {
                for (t, f) in cwise_unops {
                    let (t, f) = (id(t), id(f));
                    impls.append_all(quote! {
                        impl #t for #ty_name {
                            type Output = Self;
                            fn #f(self) -> Self {
                                Self {
                                    #(#coord: self.#coord.#f(),)*
                                }
                            }
                        }
                    });
                }
            }
        }
    }

    impls
}

fn make_mats() -> TokenStream {
    let mut tokens = TokenStream::new();
    for (&cols, &rows, &ty) in cartesian!(DIMS.iter(), DIMS.iter(), TYS.iter()) {
        let ty_ident = Ident::new(&ty.mat_name(cols, rows), Span::call_site());
        let el_ty = ty.member_ty_ident();

        tokens.append_all(quote! {
            #[derive(Copy, Clone, Debug, Default, PartialEq)]
            pub struct #ty_ident {
                buffer: [[#el_ty; #cols]; #rows],
            }
        });

        if cols == rows {
            let ty_alias = Ident::new(&ty.sq_mat_name(cols), Span::call_site());
            tokens.append_all(quote! {
                pub type #ty_alias = #ty_ident;
            });
        }

        if ty.has_arithmetic() {
            tokens.append_all(quote! {
                impl Mul<#el_ty> for #ty_ident {
                    type Output = Self;
                    fn mul(self, mul: #el_ty) -> Self {
                        let mut new = Self::default();
                        for r in 0..#rows {
                            for c in 0..#cols {
                                new.buffer[r][c] = self.buffer[r][c] * mul;
                            }
                        }
                        new
                    }
                }
                impl Add for #ty_ident {
                    type Output = #ty_ident;
                    fn add(self, other: Self) -> Self {
                        let mut new = Self::default();
                        for r in 0..#rows {
                            for c in 0..#cols {
                                new.buffer[r][c] = self.buffer[r][c] + other.buffer[r][c];
                            }
                        }
                        new
                    }
                }
                impl Sub for #ty_ident {
                    type Output = #ty_ident;
                    fn sub(self, other: Self) -> Self {
                        let mut new = Self::default();
                        for r in 0..#rows {
                            for c in 0..#cols {
                                new.buffer[r][c] = self.buffer[r][c] - other.buffer[r][c];
                            }
                        }
                        new
                    }
                }
            });
        }
        if ty.has_negation() {
            tokens.append_all(quote! {
                impl Neg for #ty_ident {
                    type Output = Self;
                    fn neg(self) -> Self {
                        let mut new = Self::default();
                        for r in 0..#rows {
                            for c in 0..#cols {
                                new.buffer[r][c] = -self.buffer[r][c];
                            }
                        }
                        new
                    }
                }
            });
        }
    }

    let mut test_idx = 0;


    // multiplications
    for (&lc, &rc, &m, &ty) in cartesian!(DIMS.iter(), DIMS.iter(), DIMS.iter(), TYS.iter()) {
        let lty = Ident::new(&ty.mat_name(lc, m), Span::call_site());
        let rty = Ident::new(&ty.mat_name(m, rc), Span::call_site());
        let oty = Ident::new(&ty.mat_name(lc, rc), Span::call_site());
        let el_ty = ty.member_ty_ident();

        if ty.has_arithmetic() {

            let test_ident = Ident::new(&format!("mul{test_idx}"), Span::call_site());

            tokens.append_all(quote! {
                impl Mul<#rty> for #lty {
                    type Output = #oty;
                    fn mul(self, rhs: #rty) -> #oty {
                        let mut result: #oty = Default::default();
                        for r in 0..#rc {
                            for c in 0..#lc {
                                let mut value: #el_ty = Default::default();
                                for m in 0..#m {
                                    value += self.buffer[m][c] * rhs.buffer[r][m];
                                }
                                result.buffer[r][c] = value;
                            }
                        }
                        result
                    }
                }

                #[test]
                fn #test_ident() {
                    let result = #lty::default() * #rty::default();
                    assert_eq!(result, #oty::default());
                }
            });

            test_idx += 1;
        }
    }
    tokens
}
