use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
use wasmtime::*;

fn interruptable_store() -> Store<()> {
    let engine = Engine::new(Config::new().interruptable(true)).unwrap();
    Store::new(&engine, ())
}

fn hugely_recursive_module(engine: &Engine) -> anyhow::Result<Module> {
    let mut wat = String::new();
    wat.push_str(
        r#"
        (import "" "" (func))
        (func (export "loop") call 2 call 2)
    "#,
    );
    for i in 0..100 {
        wat.push_str(&format!("(func call {0} call {0})\n", i + 3));
    }
    wat.push_str("(func call 0)\n");

    Module::new(engine, &wat)
}

#[test]
fn loops_interruptable() -> anyhow::Result<()> {
    let mut store = interruptable_store();
    let module = Module::new(store.engine(), r#"(func (export "loop") (loop br 0))"#)?;
    let instance = Instance::new(&mut store, &module, &[])?;
    let iloop = instance.get_typed_func::<(), (), _>(&mut store, "loop")?;
    store.interrupt_handle()?.interrupt();
    let trap = iloop.call(&mut store, ()).unwrap_err();
    assert!(trap.to_string().contains("wasm trap: interrupt"));
    Ok(())
}

#[test]
fn functions_interruptable() -> anyhow::Result<()> {
    let mut store = interruptable_store();
    let module = hugely_recursive_module(store.engine())?;
    let func = Func::wrap(&mut store, || {});
    let instance = Instance::new(&mut store, &module, &[func.into()])?;
    let iloop = instance.get_typed_func::<(), (), _>(&mut store, "loop")?;
    store.interrupt_handle()?.interrupt();
    let trap = iloop.call(&mut store, ()).unwrap_err();
    assert!(
        trap.to_string().contains("wasm trap: interrupt"),
        "{}",
        trap.to_string()
    );
    Ok(())
}

#[test]
fn loop_interrupt_from_afar() -> anyhow::Result<()> {
    // Create an instance which calls an imported function on each iteration of
    // the loop so we can count the number of loop iterations we've executed so
    // far.
    static HITS: AtomicUsize = AtomicUsize::new(0);
    let mut store = interruptable_store();
    let module = Module::new(
        store.engine(),
        r#"
            (import "" "" (func))

            (func (export "loop")
                (loop
                    call 0
                    br 0)
            )
        "#,
    )?;
    let func = Func::wrap(&mut store, || {
        HITS.fetch_add(1, SeqCst);
    });
    let instance = Instance::new(&mut store, &module, &[func.into()])?;

    // Use the instance's interrupt handle to wait for it to enter the loop long
    // enough and then we signal an interrupt happens.
    let handle = store.interrupt_handle()?;
    let thread = std::thread::spawn(move || {
        while HITS.load(SeqCst) <= 100_000 {
            // continue ...
        }
        println!("interrupting");
        handle.interrupt();
    });

    // Enter the infinitely looping function and assert that our interrupt
    // handle does indeed actually interrupt the function.
    let iloop = instance.get_typed_func::<(), (), _>(&mut store, "loop")?;
    let trap = iloop.call(&mut store, ()).unwrap_err();
    thread.join().unwrap();
    assert!(
        trap.to_string().contains("wasm trap: interrupt"),
        "bad message: {}",
        trap.to_string()
    );
    Ok(())
}

#[test]
fn function_interrupt_from_afar() -> anyhow::Result<()> {
    // Create an instance which calls an imported function on each iteration of
    // the loop so we can count the number of loop iterations we've executed so
    // far.
    static HITS: AtomicUsize = AtomicUsize::new(0);
    let mut store = interruptable_store();
    let module = hugely_recursive_module(store.engine())?;
    let func = Func::wrap(&mut store, || {
        HITS.fetch_add(1, SeqCst);
    });
    let instance = Instance::new(&mut store, &module, &[func.into()])?;

    // Use the instance's interrupt handle to wait for it to enter the loop long
    // enough and then we signal an interrupt happens.
    let handle = store.interrupt_handle()?;
    let thread = std::thread::spawn(move || {
        while HITS.load(SeqCst) <= 100_000 {
            // continue ...
        }
        handle.interrupt();
    });

    // Enter the infinitely looping function and assert that our interrupt
    // handle does indeed actually interrupt the function.
    let iloop = instance.get_typed_func::<(), (), _>(&mut store, "loop")?;
    let trap = iloop.call(&mut store, ()).unwrap_err();
    thread.join().unwrap();
    assert!(
        trap.to_string().contains("wasm trap: interrupt"),
        "bad message: {}",
        trap.to_string()
    );
    Ok(())
}