use crate::generators::Stacks;

use wasmtime::bail;

use wasmtime::*;

/// Run the given `Stacks` test case and assert that the host's view of the Wasm

/// stack matches the test case's understanding of the Wasm stack.

///

/// Returns the maximum stack depth we checked.

pub fn check_stacks(stacks: Stacks) -> usize {

    let wasm = stacks.wasm();

    crate::oracles::log_wasm(&wasm);

    let mut config = Config::new();

    config.wasm_backtrace_max_frames(stacks.limit);

    let engine = Engine::new(&config).unwrap();

    let module = Module::new(&engine, &wasm).expect("should compile okay");

    let mut linker = Linker::new(&engine);

    linker

        .func_wrap(

            "host",

            "check_stack",

            |mut caller: Caller<'_, ()>| -> Result<()> {

                let fuel = caller

                    .get_export("fuel")

                    .expect("should export `fuel`")

                    .into_global()

                    .expect("`fuel` export should be a global");

                let fuel_left = fuel.get(&mut caller).unwrap_i32();

                if fuel_left == 0 {

                    bail!(Trap::OutOfFuel);

                }

                fuel.set(&mut caller, Val::I32(fuel_left - 1)).unwrap();

                Ok(())

            },

        )

        .unwrap()

        .func_wrap(

            "host",

            "call_func",

            |mut caller: Caller<'_, ()>, f: Option<Func>| {

                let f = f.unwrap();

                let ty = f.ty(&caller);

                let params = vec![Val::I32(0); ty.params().len()];

                let mut results = vec![Val::I32(0); ty.results().len()];

                f.call(&mut caller, &params, &mut results)?;

                Ok(())

            },

        )

        .unwrap();

    let mut store = Store::new(&engine, ());

    let instance = linker

        .instantiate(&mut store, &module)

        .expect("should instantiate okay");

    let run = instance

        .get_typed_func::<(u32,), ()>(&mut store, "run")

        .expect("should export `run` function");

    let mut max_stack_depth = 0;

    for input in stacks.inputs().iter().copied() {

        log::debug!("input: {input}");

        if let Err(trap) = run.call(&mut store, (input.into(),)) {

            log::debug!("trap: {trap:?}");

            let get_stack = instance

                .get_typed_func::<(), (u32, u32)>(&mut store, "get_stack")

                .expect("should export `get_stack` function as expected");

            let (ptr, len) = get_stack

                .call(&mut store, ())

                .expect("`get_stack` should not trap");

            let memory = instance

                .get_memory(&mut store, "memory")

                .expect("should have `memory` export");

            let host_trace = match trap.downcast_ref::<WasmBacktrace>() {

                Some(bt) => bt.frames(),

                None => {

                    assert!(stacks.limit.is_none());

                    continue;

                }

            };

            let trap = trap.downcast_ref::<Trap>().unwrap();

            max_stack_depth = max_stack_depth.max(host_trace.len());

            assert_stack_matches(

                &mut store,

                memory,

                ptr,

                len,

                host_trace,

                *trap,

                stacks.limit,

            );

        }

    }

    max_stack_depth

}

/// Assert that the Wasm program's view of the stack matches the host's view.

fn assert_stack_matches(

    store: &mut impl AsContextMut,

    memory: Memory,

    ptr: u32,

    len: u32,

    host_trace: &[FrameInfo],

    trap: Trap,

    limit: Option<std::num::NonZeroUsize>,

) {

    let mut data = vec![0; len as usize];

    memory

        .read(&mut *store, ptr as usize, &mut data)

        .expect("should be in bounds");

    let mut wasm_trace = vec![];

    for entry in data.chunks(4).rev() {

        let mut bytes = [0; 4];

        bytes.copy_from_slice(entry);

        let entry = u32::from_le_bytes(bytes);

        wasm_trace.push(entry);

    }

    let trace_limit = match limit {

        Some(n) => n.get(),

        None => {

            // Backtraces are disabled; the host trace should be empty.

            assert!(host_trace.is_empty());

            return;

        }

    };

    // If the test case here trapped due to stack overflow then the host trace

    // will have one more frame than the wasm trace. The wasm didn't actually

    // get to the point of pushing onto its own trace stack where the host will

    // be able to see the exact function that triggered the stack overflow. In

    // this situation the host trace is asserted to be one larger and then the

    // top frame (first) of the host trace is discarded.

    let (host_trace, wasm_trace) = if trap == Trap::StackOverflow {

        if host_trace.len() == trace_limit {

            assert!(

                trace_limit <= wasm_trace.len() + 1,

                "Host trace size {} is larger than expected {}",

                trace_limit,

                wasm_trace.len() + 1

            );

        } else {

            assert_eq!(host_trace.len(), wasm_trace.len() + 1);

        }

        (

            &host_trace[1..],

            &wasm_trace.get(..trace_limit - 1).unwrap_or(&wasm_trace),

        )

    } else {

        (

            host_trace,

            &wasm_trace.get(..trace_limit).unwrap_or(&wasm_trace),

        )

    };

    log::debug!("Wasm thinks its stack is: {wasm_trace:?}");

    log::debug!(

        "Host thinks the stack is: {:?}",

        host_trace

            .iter()

            .map(|f| f.func_index())

            .collect::<Vec<_>>()

    );

    assert_eq!(wasm_trace.len(), host_trace.len());

    for (wasm_entry, host_entry) in wasm_trace.into_iter().zip(host_trace) {

        assert_eq!(wasm_entry, &host_entry.func_index());

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    use crate::test::gen_until_pass;

    const TARGET_STACK_DEPTH: usize = 10;

    #[test]

    fn smoke_test() {

        gen_until_pass(|stacks: Stacks, _u| {

            let max_stack_depth = check_stacks(stacks);

            Ok(max_stack_depth >= TARGET_STACK_DEPTH)

        });

    }

}