use bevy_ecs::prelude::*;
use wasmtime::*;

#[derive(Clone)]
struct CustomMemory {
    /// This pointer is not owned by CustomMemory!
    ptr: usize,
    size: usize,
}

unsafe impl LinearMemory for CustomMemory {
    fn byte_size(&self) -> usize {
        self.size
    }

    fn byte_capacity(&self) -> usize {
        self.size
    }

    fn grow_to(&mut self, _new_size: usize) -> wasmtime::Result<()> {
        Err(anyhow::anyhow!("It is not possible to grow this memory"))
    }

    fn as_ptr(&self) -> *mut u8 {
        self.ptr as *mut u8
    }
}

struct CustomMemoryCreator(CustomMemory);

unsafe impl MemoryCreator for CustomMemoryCreator {
    fn new_memory(
        &self,
        ty: MemoryType,
        _minimum: usize,
        _maximum: Option<usize>,
        reserved_size: Option<usize>,
        guard_size: usize,
    ) -> wasmtime::Result<Box<dyn LinearMemory>, String> {
        assert_eq!(guard_size, 0);
        assert_eq!(reserved_size, Some(0));
        assert!(!ty.is_64());
        assert_eq!(ty.page_size(), 1);
        Ok(Box::new(self.0.clone()))
    }
}

// Note: this type is technically aligned = 1
#[derive(Resource)]
struct Res(u8);

fn main() {
    let mut world = World::new();
    world.insert_resource(Res(12));

    {
        let component_id = world.components().resource_id::<Res>().unwrap();
        let ptr = world
            .storages()
            .resources
            .get(component_id)
            .unwrap()
            .get_data()
            .unwrap()
            .as_ptr();

        // Note: Res is not aligned=4, so isn't guaranteed to be aligned correctly for wasm memory
        // yet this always passes because my 64 bit machine aligns allocations to 16
        assert_eq!(ptr as usize % 4, 0);

        // Repr should be transparent
        assert_eq!(unsafe { *ptr.cast::<u8>() }, 12);

        let mem_creator = std::sync::Arc::new(CustomMemoryCreator(CustomMemory {
            ptr: ptr as _,
            size: size_of::<Res>(),
        }));

        let engine = Engine::new(
            Config::new()
                .wasm_custom_page_sizes(true)
                .with_host_memory(mem_creator.clone())
                .memory_reservation(0)
                .memory_guard_size(0),
        )
        .unwrap();

        let bytes = br#"
            (module
                (memory (export "memory") 1 (pagesize 1))

                (func (export "run")
                    ;; Attempting a regular i32.store and i32.load would trap
                    (i32.store8
                        (i32.const 0)
                        (i32.add
                            (i32.load8_u (i32.const 0))
                            (i32.const 123)
                        )
                    )
                )
            )
        "#;
        let module = Module::new(&engine, bytes).unwrap();

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

        let instance = Instance::new(&mut store, &module, &[]).unwrap();

        let run = instance
            .get_typed_func::<(), ()>(&mut store, "run")
            .unwrap();

        run.call(&mut store, ()).unwrap();
    }

    assert_eq!(world.get_resource::<Res>().unwrap().0, 123 + 12);
}