use wasmtime::component::{Component, Linker, TaskGroupHook, TaskGroupId};
use wasmtime::{Config, Engine, Result, Store};

struct Hook;

impl TaskGroupHook for Hook {
    fn handle_start(&mut self, id: TaskGroupId) -> Result<()> {
        println!("start  {id:?}");
        Ok(())
    }
    fn handle_enter(&mut self, id: TaskGroupId) -> Result<()> {
        println!("enter  {id:?}");
        Ok(())
    }
    fn handle_exit(&mut self, id: TaskGroupId) -> Result<()> {
        println!("exit   {id:?}");
        Ok(())
    }
    fn handle_finish(&mut self, id: TaskGroupId) -> Result<()> {
        println!("finish {id:?}");
        Ok(())
    }
}

const WAT: &str = r#"
(component
  (core module $m
    (import "" "task.return" (func $task.return))
    (import "" "backpressure.inc" (func $backpressure.inc))
    ;; `bg`: turn on backpressure, return a result, then keep running in the
    ;; background by yielding; the next callback invocation traps.
    (func (export "bg") (result i32)
      (call $backpressure.inc)
      (call $task.return)
      (i32.const 1 (; YIELD ;)))
    (func (export "bg-cb") (param i32 i32 i32) (result i32)
      unreachable)
    ;; `foo`: never gets to run because of the backpressure.
    (func (export "foo") (result i32)
      (call $task.return)
      (i32.const 0 (; EXIT ;)))
    (func (export "foo-cb") (param i32 i32 i32) (result i32)
      unreachable)
  )
  (core func $task.return (canon task.return))
  (core func $backpressure.inc (canon backpressure.inc))
  (core instance $i (instantiate $m (with "" (instance
    (export "task.return" (func $task.return))
    (export "backpressure.inc" (func $backpressure.inc))))))
  (func (export "bg") async
    (canon lift (core func $i "bg") async (callback (core func $i "bg-cb"))))
  (func (export "foo") async
    (canon lift (core func $i "foo") async (callback (core func $i "foo-cb"))))
)
"#;

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<()> {
    let mut config = Config::new();
    config.wasm_component_model_async(true);
    let engine = Engine::new(&config)?;
    let component = Component::new(&engine, WAT)?;
    let mut store = Store::new(&engine, ());
    if std::env::var_os("NO_HOOK").is_none() {
        store.task_group_hook(Hook);
    }
    let linker = Linker::new(&engine);
    let instance = linker.instantiate_async(&mut store, &component).await?;
    let bg = instance.get_typed_func::<(), ()>(&mut store, "bg")?;
    let foo = instance.get_typed_func::<(), ()>(&mut store, "foo")?;

    println!("calling bg");
    bg.call_async(&mut store, ()).await?;
    println!("bg returned; calling foo");
    let r = foo.call_async(&mut store, ()).await;
    println!("foo result: {r:?}");
    Ok(())
}