//! Test to assert that resumable call feature works as intended.

use core::slice;

use wasmi::{
    Config,
    Engine,
    Error,
    Extern,
    Func,
    FuncType,
    Linker,
    Module,
    ResumableCall,
    ResumableInvocation,
    Store,
    TypedFunc,
    TypedResumableCall,
    TypedResumableInvocation,
    Value,
};
use wasmi_core::{Trap, TrapCode, ValueType};

fn test_setup() -> (Store<()>, Linker<()>) {
    let mut config = Config::default();
    config.wasm_tail_call(true);
    let engine = Engine::new(&config);
    let store = Store::new(&engine, ());
    let linker = <Linker<()>>::new(&engine);
    (store, linker)
}

fn resumable_call_smoldot_common(wasm: &str) -> (Store<()>, TypedFunc<(), i32>) {
    let (mut store, mut linker) = test_setup();
    // The important part about this test is that this
    // host function has more results than parameters.
    let host_fn = Func::new(
        &mut store,
        FuncType::new([], [ValueType::I32]),
        |_caller, _inputs, _output| Err(Trap::i32_exit(100)),
    );
    linker.define("env", "host_fn", host_fn).unwrap();
    // The Wasm defines a single function that calls the
    // host function, returns 10 if the output is 0 and
    // returns 20 otherwise.
    let wasm = wat::parse_str(wasm).unwrap();
    let module = Module::new(store.engine(), &mut &wasm[..]).unwrap();
    let instance = linker
        .instantiate(&mut store, &module)
        .unwrap()
        .start(&mut store)
        .unwrap();
    let wasm_fn = instance.get_typed_func::<(), i32>(&store, "test").unwrap();
    (store, wasm_fn)
}

pub trait UnwrapResumable {
    type Results;

    fn unwrap_resumable(self) -> TypedResumableInvocation<Self::Results>;
}

impl<Results> UnwrapResumable for Result<TypedResumableCall<Results>, Trap> {
    type Results = Results;

    fn unwrap_resumable(self) -> TypedResumableInvocation<Self::Results> {
        match self.unwrap() {
            TypedResumableCall::Resumable(invocation) => invocation,
            TypedResumableCall::Finished(_) => panic!("expected TypedResumableCall::Resumable"),
        }
    }
}

#[test]
fn resumable_call_smoldot_01() {
    let (mut store, wasm_fn) = resumable_call_smoldot_common(
        r#"
        (module
            (import "env" "host_fn" (func $host_fn (result i32)))
            (func (export "test") (result i32)
                (call $host_fn)
            )
        )
        "#,
    );
    let invocation = wasm_fn.call_resumable(&mut store, ()).unwrap_resumable();
    match invocation.resume(&mut store, &[Value::I32(42)]).unwrap() {
        TypedResumableCall::Finished(result) => assert_eq!(result, 42),
        TypedResumableCall::Resumable(_) => panic!("expected TypeResumableCall::Finished"),
    }
}

#[test]
fn resumable_call_smoldot_tail_01() {
    let (mut store, wasm_fn) = resumable_call_smoldot_common(
        r#"
        (module
            (import "env" "host_fn" (func $host_fn (result i32)))
            (func (export "test") (result i32)
                (return_call $host_fn)
            )
        )
        "#,
    );
    assert_eq!(
        wasm_fn
            .call_resumable(&mut store, ())
            .unwrap_err()
            .i32_exit_status(),
        Some(100),
    );
}

#[test]
fn resumable_call_smoldot_tail_02() {
    let (mut store, wasm_fn) = resumable_call_smoldot_common(
        r#"
        (module
            (import "env" "host_fn" (func $host (result i32)))
            (func $wasm (result i32)
                (return_call $host)
            )
            (func (export "test") (result i32)
                (call $wasm)
            )
        )
        "#,
    );
    let invocation = wasm_fn.call_resumable(&mut store, ()).unwrap_resumable();
    match invocation.resume(&mut store, &[Value::I32(42)]).unwrap() {
        TypedResumableCall::Finished(result) => assert_eq!(result, 42),
        TypedResumableCall::Resumable(_) => panic!("expected TypeResumableCall::Finished"),
    }
}

#[test]
fn resumable_call_smoldot_02() {
    let (mut store, wasm_fn) = resumable_call_smoldot_common(
        r#"
        (module
            (import "env" "host_fn" (func $host_fn (result i32)))
            (func (export "test") (result i32)
                (if (result i32) (i32.ne (call $host_fn) (i32.const 0))
                    (then
                        (i32.const 11) ;; EXPECTED
                    )
                    (else
                        (i32.const 22) ;; FAILURE
                    )
                )
            )
        )
        "#,
    );
    let invocation = wasm_fn.call_resumable(&mut store, ()).unwrap_resumable();
    match invocation.resume(&mut store, &[Value::I32(42)]).unwrap() {
        TypedResumableCall::Finished(result) => assert_eq!(result, 11),
        TypedResumableCall::Resumable(_) => panic!("expected TypeResumableCall::Finished"),
    }
}

#[test]
fn resumable_call_host() {
    let (mut store, _linker) = test_setup();
    let host_fn = Func::wrap(&mut store, || -> Result<(), Trap> {
        Err(Trap::i32_exit(100))
    });
    // Even though the called host function traps we expect a normal error
    // since the host function is the root function of the call and therefore
    // it would not make sense to resume it.
    let error = host_fn
        .call_resumable(&mut store, &[], &mut [])
        .unwrap_err();
    match error {
        Error::Trap(trap) => {
            assert_eq!(trap.i32_exit_status(), Some(100));
        }
        _ => panic!("expected Wasm trap"),
    }
    // The same test for `TypedFunc`:
    let trap = host_fn
        .typed::<(), ()>(&store)
        .unwrap()
        .call_resumable(&mut store, ())
        .unwrap_err();
    assert_eq!(trap.i32_exit_status(), Some(100));
}

#[test]
fn resumable_call() {
    let (mut store, mut linker) = test_setup();
    let host_fn = Func::wrap(&mut store, |input: i32| -> Result<i32, Trap> {
        match input {
            1 => Err(Trap::i32_exit(10)),
            2 => Err(Trap::i32_exit(20)),
            n => Ok(n + 1),
        }
    });
    linker.define("env", "host_fn", host_fn).unwrap();
    let wasm = wat::parse_str(
        r#"
        (module
            (import "env" "host_fn" (func $host_fn (param i32) (result i32)))
            (func (export "wasm_fn") (param $wasm_trap i32) (result i32)
                (local $i i32)
                (local.set $i (i32.const 0))
                (local.set $i (call $host_fn (local.get $i))) ;; Ok
                (local.set $i (call $host_fn (local.get $i))) ;; Trap::i32_exit(1)
                (local.set $i (call $host_fn (local.get $i))) ;; Trap::i32_exit(2)
                (local.set $i (call $host_fn (local.get $i))) ;; Ok
                (if (i32.eq (local.get $wasm_trap) (i32.const 1))
                    (then unreachable)                        ;; trap in Wasm if $wasm_trap == 1
                )
                (local.get $i)                                ;; return i == 4
            )
        )
        "#,
    )
    .unwrap();

    let module = Module::new(store.engine(), &mut &wasm[..]).unwrap();
    let instance = linker
        .instantiate(&mut store, &module)
        .unwrap()
        .start(&mut store)
        .unwrap();
    let wasm_fn = instance
        .get_export(&store, "wasm_fn")
        .and_then(Extern::into_func)
        .unwrap();

    run_test(wasm_fn, &mut store, false);
    run_test(wasm_fn, &mut store, true);
    run_test_typed(wasm_fn, &mut store, false);
    run_test_typed(wasm_fn, &mut store, true);
}

trait AssertResumable {
    type Results;
    type Invocation;

    fn assert_resumable(
        self,
        store: &Store<()>,
        exit_status: i32,
        host_results: &[ValueType],
    ) -> Self::Invocation;
    fn assert_finish(self) -> Self::Results;
}

impl AssertResumable for ResumableCall {
    type Results = ();
    type Invocation = ResumableInvocation;

    fn assert_resumable(
        self,
        store: &Store<()>,
        exit_status: i32,
        host_results: &[ValueType],
    ) -> Self::Invocation {
        match self {
            Self::Resumable(invocation) => {
                assert_eq!(invocation.host_error().i32_exit_status(), Some(exit_status));
                assert_eq!(invocation.host_func().ty(store).results(), host_results,);
                invocation
            }
            Self::Finished => panic!("expected host function trap with exit code 10"),
        }
    }

    fn assert_finish(self) -> Self::Results {
        match self {
            Self::Finished => (),
            Self::Resumable(_) => panic!("expected the resumable call to finish"),
        }
    }
}

fn run_test(wasm_fn: Func, mut store: &mut Store<()>, wasm_trap: bool) {
    let mut results = Value::I32(0);
    let invocation = wasm_fn
        .call_resumable(
            &mut store,
            &[Value::I32(wasm_trap as i32)],
            slice::from_mut(&mut results),
        )
        .unwrap()
        .assert_resumable(store, 10, &[ValueType::I32]);
    let invocation = invocation
        .resume(&mut store, &[Value::I32(2)], slice::from_mut(&mut results))
        .unwrap()
        .assert_resumable(store, 20, &[ValueType::I32]);
    let call = invocation.resume(&mut store, &[Value::I32(3)], slice::from_mut(&mut results));
    if wasm_trap {
        match call.unwrap_err() {
            Error::Trap(trap) => {
                assert!(matches!(
                    trap.trap_code(),
                    Some(TrapCode::UnreachableCodeReached)
                ));
            }
            _ => panic!("expected Wasm trap"),
        }
    } else {
        call.unwrap().assert_finish();
        assert_eq!(results.i32(), Some(4));
    }
}

impl<Results> AssertResumable for TypedResumableCall<Results> {
    type Results = Results;
    type Invocation = TypedResumableInvocation<Results>;

    fn assert_resumable(
        self,
        store: &Store<()>,
        exit_status: i32,
        host_results: &[ValueType],
    ) -> Self::Invocation {
        match self {
            Self::Resumable(invocation) => {
                assert_eq!(invocation.host_error().i32_exit_status(), Some(exit_status));
                assert_eq!(invocation.host_func().ty(store).results(), host_results,);
                invocation
            }
            Self::Finished(_) => panic!("expected host function trap with exit code 10"),
        }
    }

    fn assert_finish(self) -> Self::Results {
        match self {
            Self::Finished(results) => results,
            Self::Resumable(_) => panic!("expected the resumable call to finish"),
        }
    }
}

fn run_test_typed(wasm_fn: Func, mut store: &mut Store<()>, wasm_trap: bool) {
    let invocation = wasm_fn
        .typed::<i32, i32>(&store)
        .unwrap()
        .call_resumable(&mut store, wasm_trap as i32)
        .unwrap()
        .assert_resumable(store, 10, &[ValueType::I32]);
    let invocation = invocation
        .resume(&mut store, &[Value::I32(2)])
        .unwrap()
        .assert_resumable(store, 20, &[ValueType::I32]);
    let call = invocation.resume(&mut store, &[Value::I32(3)]);
    if wasm_trap {
        match call.unwrap_err() {
            Error::Trap(trap) => {
                assert!(matches!(
                    trap.trap_code(),
                    Some(TrapCode::UnreachableCodeReached)
                ));
            }
            _ => panic!("expected Wasm trap"),
        }
    } else {
        assert_eq!(call.unwrap().assert_finish(), 4);
    }
}
