use crate::module_macro;
use proc_macro2::TokenStream;
use quote::quote;
use syn::visit::Visit;

#[derive(Debug)]
struct CompileErrorFinder {
    error_msg: &'static str,
    found: bool,
    found_errors: Vec<String>,
}

impl syn::visit::Visit<'_> for CompileErrorFinder {
    fn visit_macro(&mut self, mac: &syn::Macro) {
        if let Some(segment) = mac.path.segments.last() {
            if segment.ident == "compile_error" {
                if let Ok(lit_str) = syn::parse2::<syn::LitStr>(mac.tokens.clone()) {
                    let lowercase = lit_str.value().to_lowercase();
                    if lowercase.contains(self.error_msg) {
                        self.found = true;
                    }
                    self.found_errors.push(lowercase);
                }
            }
        }
    }
}

fn assert_has_compile_error(error_msg: &'static str, token_stream: TokenStream) {
    let file = syn::parse2::<syn::File>(token_stream).unwrap();
    let mut finder = CompileErrorFinder {
        error_msg,
        found: false,
        found_errors: vec![],
    };
    finder.visit_file(&file);
    assert!(
        finder.found,
        "{error_msg} was not found. Found_errors: {:?}",
        finder.found_errors
    );
}

macro_rules! module_fails_with_error_message {
    ($error_msg:expr, $($code:tt)+) => {
        {
            let code = quote! {$($code)+};
            let result = module_macro(code);
            assert_has_compile_error($error_msg, result);
        }
    };
}
macro_rules! module_fails_with_error_messages {
    ($($error_msg:expr)+, $($code:tt)+) => {
        {
            let code = quote! {$($code)+};
            let result = module_macro(code);
            $(assert_has_compile_error($error_msg, result.clone());)+
        }
    };
}

#[test]
fn test_while() {
    module_fails_with_error_messages!(
        "`while`" "unsupported" "`return`",
        mod blabla {
            #[instruction]
            fn func() -> u8 {
                while true {}
                let arr = [1, 2, 3];
                let s = &arr[1..3];
                return s[0]
            }
        }
    )
}
#[test]
fn test_recursion() {
    module_fails_with_error_message!(
        "recursion",
        mod blabla {
            #[instruction]
            fn func() -> u8 {
                func()
            }
        }
    )
}

#[test]
fn test_if_reveal() {
    module_fails_with_error_messages!(
        "reveal" "if location",
        mod blabla {
            #[instruction]
            fn func(a: u8, b: bool) -> u8 {
                if b {
                    a.reveal()
                } else {
                    a
                }
            }
        }
    );
    module_fails_with_error_messages!(
        "from_arcis" "if location",
        mod blabla {
            #[instruction]
            fn func(b: bool, c0: Shared, c1: Shared, val: u8) -> Enc<Shared, u8> {
                if b {
                    c0.from_arcis(val)
                } else {
                    c1.from_arcis(val)
                }
            }
        }
    );
    module_fails_with_error_messages!(
        "from_arcis" "if location",
        mod blabla {
            #[instruction]
            fn func(b: bool, owner: Mxe, val0: u8, val1: u8) -> Enc<Mxe, u8> {
                if b {
                    owner.from_arcis(val0)
                } else {
                    owner.from_arcis(val1)
                }
            }
        }
    );
}

#[test]
fn test_empty_module() {
    module_fails_with_error_message!(
        "#[instruction]",
        mod blabla {
            fn func() -> u8 {
                0
            }
        }
    )
}
#[test]
fn test_instruction_in_sub_module() {
    module_fails_with_error_message!(
        "#[encrypted] module's top level",
        mod blabla {
            mod example {
                #[instruction]
                fn func() -> u8 {
                    0
                }
            }
            #[instruction]
            fn main() -> u8 {
                0
            }
        }
    )
}

#[test]
fn test_recursive_const() {
    // rust will catch it but we don't want an infinite loop.
    module_fails_with_error_message!(
        "recursive",
        mod blabla {
            const C_0: usize = C_1;
            const C_1: usize = C_0;
            #[instruction]
            fn func() -> usize {
                C_0
            }
        }
    )
}

#[test]
fn test_assign_conversion() {
    module_fails_with_error_message!(
        "conversion",
        mod blabla {
            #[instruction]
            fn see_what_happens() -> [u8; 2] {
                let mut arr = [0u8, 1u8];
                arr[0] = 0.0;
                arr
            }
        }
    )
}

#[test]
fn test_take_var() {
    module_fails_with_error_message!(
        "const",
        mod blabla {
            #[instruction]
            fn take_var(arr: [u8; 1], b: bool) -> u8 {
                let mut res = 0;
                for i in arr.into_iter().take(b as usize) {
                    res += i;
                }
                res
            }
        }
    )
}

#[test]
fn test_side_effects_macro() {
    module_fails_with_error_messages!(
        "side effects" "`println!`",
        mod blabla {
            fn post_increment(val: &mut u32) -> u32 {
                let old_val = *val;
                *val += 1;
                old_val
            }
            #[instruction]
            fn foo(mut n: u32) -> u32 {
                println!("{}", post_increment(&mut n));
                n
            }
        }
    )
}
#[test]
fn test_string_support() {
    module_fails_with_error_messages!(
        "b\"hello_world\"",
        mod blabla {
            #[instruction]
            fn foo(n: u32) -> u32 {
                let a = "hello_world";
                n
            }
        }
    )
}
#[test]
fn test_neg_u32() {
    module_fails_with_error_messages!(
        "unsigned" "opposite",
        mod blabla {
            #[instruction]
            fn foo(n: u32) -> u32 {
                -n
            }
        }
    )
}
#[test]
fn test_fn_comparison() {
    module_fails_with_error_messages!(
        "unimplemented" "`fn(..) -> _`",
        mod blabla {
            #[instruction]
            fn foo() -> bool {
                (|| {}) == (|| {})
            }
        }
    )
}
#[test]
fn test_arr_field() {
    module_fails_with_error_messages!(
        ".0" "tuple",
        mod blabla {
            #[instruction]
            fn foo() -> bool {
                [true].0
            }
        }
    )
}
#[test]
fn test_deref_conversion() {
    module_fails_with_error_messages!(
        "conversion",
        mod blabla {
            #[instruction]
            fn foo(b: &bool) -> bool {
                b as bool
            }
        }
    )
}
#[test]
fn test_index_u8() {
    module_fails_with_error_messages!(
        "`usize`" "`isize`" "index",
        mod blabla {
            #[instruction]
            fn foo(b: u8, arr: [bool; 256]) -> bool {
                arr[b]
            }
        }
    )
}
#[test]
fn test_size_isize() {
    module_fails_with_error_messages!(
        "`usize`",
        mod blabla {
            #[instruction]
            fn foo() -> bool {
                [true; 2isize][0]
            }
        }
    )
}
#[test]
fn test_unknown_input_type() {
    module_fails_with_error_messages!(
        "unknown type" "input type",
        mod blabla {
            #[instruction]
            fn foo(a: String) -> bool {
                true
            }
        }
    )
}
#[test]
fn test_unknown_return_type() {
    module_fails_with_error_messages!(
        "unknown type" "return type",
        mod blabla {
            #[instruction]
            fn foo(a: bool) -> Vec {
            }
        }
    )
}
#[test]
fn test_fn_input_type() {
    module_fails_with_error_messages!(
        "function" "input type",
        mod blabla {
            #[instruction]
            fn foo(a: fn() -> ()) -> bool {
                true
            }
        }
    )
}
#[test]
fn test_fn_return_type() {
    module_fails_with_error_messages!(
        "function" "return type",
        mod blabla {
            #[instruction]
            fn foo(a: bool) -> fn() -> () {
                || {}
            }
        }
    )
}
#[test]
fn test_fn_input_arg() {
    module_fails_with_error_messages!(
        "could not find",
        mod blabla {
            #[instruction]
            fn foo(a: u8) -> u8 {
                a
            }
            #[instruction]
            fn bar(b: u8) -> u8 {
                a + b
            }
        }
    )
}

#[test]
fn test_static_panic() {
    module_fails_with_error_messages!(
        "length" "arcis_static_panic",
        mod blabla {
            #[instruction]
            fn foo() -> ArcisX25519Pubkey {
                ArcisX25519Pubkey::from_uint8(b"0000000000000000000000000000000000000000000000000")
            }
        }
    )
}

#[test]
fn test_or_with_ptr() {
    module_fails_with_error_messages!(
        "compatible" "type",
        mod blabla {
            #[instruction]
            fn foo(a: u128) -> u8 {
                (true || & a).0
            }
        }
    )
}

#[test]
fn test_div_zero() {
    module_fails_with_error_messages!(
        "division by zero",
        mod blabla {
            #[instruction]
            fn foo(a: u8) -> u8 {
                a / 0
            }
        }
    )
}

#[test]
fn test_rem_zero() {
    module_fails_with_error_messages!(
        "division by zero",
        mod blabla {
            #[instruction]
            fn foo(a: u8) -> u8 {
                a % 0
            }
        }
    )
}
#[test]
fn test_const_binding() {
    module_fails_with_error_messages!(
        "find" "`a`",
        mod blabla {
            #[instruction]
            fn foo(a: u8) -> u8 {
                const { a }
            }
        }
    )
}

#[ignore = "test does not work"]
#[test]
fn test_backtrace() {
    module_fails_with_error_messages!(
        "backtrace",
        mod blabla {
            #[instruction]
            fn foo(n: u32) -> ArcisField {
                ArcisField::from(n)
            }
        }
    )
}

#[ignore = "benchmark"]
#[test]
fn profile_module() {
    let module = quote! {
        mod decryption {
            use arcis::*;

            struct MyStruct2 {
                a: Enc<Mxe, u8>,
                b: Enc<Shared, u8>,
            }
            #[derive(Clone, Copy)]
            struct MyStruct {
                a: u8,
            }
            #[instruction]
            fn decryption(
                s: MyStruct2,
                c1: Enc<Shared, MyStruct>,
                c2: Enc<Mxe, MyStruct>,
            ) -> [u8; 4] {
                [s.a.to_arcis(), s.b.to_arcis(), c1.to_arcis().a, c2.to_arcis().a]
            }

            #[instruction]
            fn all(
                s0: MyStruct,
                s_c1: Enc<Shared, MyStruct>,
                s_c2: Enc<Mxe, MyStruct>,
                c: Shared,
            ) -> (
                Enc<Shared, MyStruct>,
                Enc<Shared, MyStruct>,
                Enc<Mxe, MyStruct>,
            ) {
                let s1 = s_c1.to_arcis();
                let s2 = s_c2.to_arcis();

                let r0 = s_c1.owner.from_arcis(s0);
                let r1 = c.from_arcis(s2);
                let r2 = s_c2.owner.from_arcis(s1);
                (r0, r1, r2)
            }
        }
    };
    module_macro(module);
}
