pub use canary_types::*;
use parking_lot::Mutex;

/// Low-level script API callbacks.
pub trait ScriptAbi {
    fn start_draw(&self);
    fn draw_indexed(&self, vertices: &[MeshVertex], indices: &[MeshIndex]);
    fn with_draw_commands(&self, f: impl FnOnce(&[DrawCommand]));
}

pub trait ScriptInstance {
    fn update(&mut self, dt: f32);
    fn draw(&mut self, f: impl FnOnce(&[DrawCommand]));
    fn on_cursor_event(&mut self, kind: CursorEventKind, at: Vec2);
}

#[non_exhaustive]
pub enum DrawCommand {
    Mesh {
        vertices: Vec<MeshVertex>,
        indices: Vec<MeshIndex>,
    },
}

pub struct WasmtimeRuntime {
    engine: wasmtime::Engine,
}

impl WasmtimeRuntime {
    pub fn new() -> anyhow::Result<Self> {
        let mut config = wasmtime::Config::new();
        config.wasm_simd(true);
        config.wasm_bulk_memory(true);
        config.cranelift_opt_level(wasmtime::OptLevel::Speed);

        let engine = wasmtime::Engine::new(&config)?;

        Ok(Self { engine })
    }

    pub fn load_module<T: ScriptAbi>(
        &self,
        abi: T,
        module: &[u8],
    ) -> anyhow::Result<WasmtimeScript<T>> {
        let module = wasmtime::Module::new(&self.engine, module)?;
        let mut store = wasmtime::Store::new(&self.engine, abi);
        let mut linker = wasmtime::Linker::new(&self.engine);
        WasmtimeScript::link(&mut linker)?;
        let instance = linker.instantiate(&mut store, &module)?;
        let update = instance.get_typed_func(&mut store, "update")?;
        let draw = instance.get_typed_func(&mut store, "draw")?;
        let on_cursor_event = instance.get_typed_func(&mut store, "on_cursor_event")?;

        let bind_panel = instance.get_typed_func::<u32, u32, _>(&mut store, "bind_panel")?;
        bind_panel.call(&mut store, 0u32)?;

        Ok(WasmtimeScript {
            store,
            update,
            draw,
            on_cursor_event,
        })
    }
}

pub struct WasmtimeScript<T> {
    store: wasmtime::Store<T>,
    update: wasmtime::TypedFunc<f32, ()>,
    draw: wasmtime::TypedFunc<(), ()>,
    on_cursor_event: wasmtime::TypedFunc<(u32, f32, f32), ()>,
}

impl<T: ScriptAbi> WasmtimeScript<T> {
    pub fn link(linker: &mut wasmtime::Linker<T>) -> anyhow::Result<()> {
        let module = "env";

        linker.func_wrap(
            module,
            "draw_indexed",
            |mut caller: wasmtime::Caller<'_, T>,
             vertices_ptr: u32,
             vertices_num: u32,
             indices_ptr: u32,
             indices_num: u32| {
                let vertices = Self::get_memory_slice(&mut caller, vertices_ptr, vertices_num);
                let indices = Self::get_memory_slice(&mut caller, indices_ptr, indices_num);
                caller.data().draw_indexed(vertices, indices);
            },
        )?;

        linker.func_wrap(
            module,
            "UiPanel_drawTriangle",
            |caller: wasmtime::Caller<'_, T>,
             _panel: u32,
             x1: f32,
             y1: f32,
             x2: f32,
             y2: f32,
             x3: f32,
             y3: f32,
             r: f32,
             g: f32,
             b: f32,
             a: f32| {
                let v1 = Vec2 { x: x1, y: y1 };
                let v2 = Vec2 { x: x2, y: y2 };
                let v3 = Vec2 { x: x3, y: y3 };
                let color = Color { r, g, b, a };

                let vertices = [
                    MeshVertex {
                        position: v1,
                        color,
                    },
                    MeshVertex {
                        position: v2,
                        color,
                    },
                    MeshVertex {
                        position: v3,
                        color,
                    },
                ];

                let indices = [0, 1, 2];

                caller.data().draw_indexed(&vertices, &indices);
            },
        )?;

        Ok(())
    }

    fn get_memory_slice<D: bytemuck::Pod>(
        caller: &mut wasmtime::Caller<'_, T>,
        ptr: u32,
        num: u32,
    ) -> &'static [D] {
        let len = num as usize * std::mem::size_of::<D>();
        let bytes = Self::get_memory_slice_bytes(caller, ptr as usize, len);
        bytemuck::cast_slice(bytes)
    }

    fn get_memory_slice_bytes(
        caller: &mut wasmtime::Caller<'_, T>,
        ptr: usize,
        len: usize,
    ) -> &'static [u8] {
        let memory = caller.get_export("memory").unwrap().into_memory().unwrap();
        if ptr + len > memory.data_size(&caller) {
            panic!("Attempted wasm memory read is out-of-bounds!");
        }

        unsafe {
            let ptr = memory.data_ptr(caller).add(ptr);
            std::slice::from_raw_parts(ptr, len)
        }
    }
}

impl<T: ScriptAbi> ScriptInstance for WasmtimeScript<T> {
    fn update(&mut self, dt: f32) {
        self.update.call(&mut self.store, dt).unwrap();
    }

    fn draw(&mut self, f: impl FnOnce(&[DrawCommand])) {
        self.store.data().start_draw();
        self.draw.call(&mut self.store, ()).unwrap();
        self.store.data().with_draw_commands(f);
    }

    fn on_cursor_event(&mut self, kind: CursorEventKind, at: Vec2) {
        self.on_cursor_event
            .call(&mut self.store, (kind as u32, at.x, at.y))
            .unwrap();
    }
}

#[derive(Default)]
pub struct SimpleScriptAbi {
    draw_cmds: Mutex<Vec<DrawCommand>>,
}

impl ScriptAbi for SimpleScriptAbi {
    fn start_draw(&self) {
        self.draw_cmds.lock().clear();
    }

    fn draw_indexed(&self, vertices: &[MeshVertex], indices: &[MeshIndex]) {
        self.draw_cmds.lock().push(DrawCommand::Mesh {
            vertices: vertices.to_vec(),
            indices: indices.to_vec(),
        })
    }

    fn with_draw_commands(&self, f: impl FnOnce(&[DrawCommand])) {
        f(self.draw_cmds.lock().as_slice());
    }
}