pub use canary_types::*;
use parking_lot::{Mutex, RwLock};
use slab::Slab;
use std::collections::HashMap;
use std::sync::Arc;

pub mod text;

/// Low-level script API callbacks.
///
/// If you're a casual user of Canary the struct you're looking for is
/// [ScriptAbiImpl]. This trait exists to help with making mocks for testing.
pub trait ScriptAbi {
    fn start_draw(&self);
    fn draw_indexed(&self, vertices: &[MeshVertex], indices: &[MeshIndex]);
    fn draw_text_layout(&self, id: u32, offset: Vec2, scale: f32, color: Color);
    fn with_draw_commands(&self, f: impl FnOnce(&[DrawCommand]));

    fn font_load(&self, family: &str) -> u32;

    fn text_layout_new(&self, font_id: u32, text: &str) -> u32;
    fn text_layout_delete(&self, id: u32);
    fn text_layout_get_bounds(&self, id: u32, rect: &mut Rect);

    fn message_new(&self, data: Vec<u8>) -> u32;
    fn message_free(&self, id: u32);
    fn message_get_len(&self, id: u32) -> u32;
    fn message_get_data(&self, id: u32, dst: &mut [u8]);
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct PanelId(pub(crate) usize);

pub trait ScriptInstance {
    fn bind_panel(&mut self, msg: Vec<u8>) -> PanelId;
    fn update(&mut self, panel: PanelId, dt: f32);
    fn draw(&mut self, panel: PanelId, f: impl FnOnce(&[DrawCommand]));
    fn on_cursor_event(&mut self, panel: PanelId, kind: CursorEventKind, at: Vec2);
    fn on_message(&mut self, panel: PanelId, msg: Vec<u8>);
}

#[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 panel_datas = Default::default();
        let mut linker = wasmtime::Linker::new(&self.engine);
        WasmtimeScript::link(&mut linker)?;
        let instance = linker.instantiate(&mut store, &module)?;
        let bind_panel = instance.get_typed_func(&mut store, "bind_panel")?;
        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 on_message = instance.get_typed_func(&mut store, "on_message")?;

        Ok(WasmtimeScript {
            store,
            panel_datas,
            bind_panel,
            update,
            draw,
            on_cursor_event,
            on_message,
        })
    }
}

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

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,
            "draw_text_layout",
            |caller: wasmtime::Caller<'_, T>,
             id: u32,
             xoff: f32,
             yoff: f32,
             scale: f32,
             r: f32,
             g: f32,
             b: f32,
             a: f32| {
                let offset = Vec2 { x: xoff, y: yoff };
                let color = Color { r, g, b, a };
                caller.data().draw_text_layout(id, offset, scale, color);
            },
        )?;

        linker.func_wrap(
            module,
            "font_load",
            |mut caller: wasmtime::Caller<'_, T>, family_ptr: u32, family_len: u32| {
                let family = Self::get_memory_slice_str(&mut caller, family_ptr, family_len);
                caller.data().font_load(family)
            },
        )?;

        linker.func_wrap(
            module,
            "text_layout_new",
            |mut caller: wasmtime::Caller<'_, T>, font_id: u32, text_ptr: u32, text_len: u32| {
                let text = Self::get_memory_slice_str(&mut caller, text_ptr, text_len);
                caller.data().text_layout_new(font_id, text)
            },
        )?;

        linker.func_wrap(
            module,
            "text_layout_delete",
            |caller: wasmtime::Caller<'_, T>, id: u32| caller.data().text_layout_delete(id),
        )?;

        linker.func_wrap(
            module,
            "text_layout_get_bounds",
            |mut caller: wasmtime::Caller<'_, T>, id: u32, rect_ptr: u32| {
                let rect: &mut Rect = Self::get_memory_ref(&mut caller, rect_ptr);
                caller.data().text_layout_get_bounds(id, rect);
            },
        )?;

        linker.func_wrap(
            module,
            "message_get_len",
            |caller: wasmtime::Caller<'_, T>, id: u32| -> u32 {
                caller.data().message_get_len(id)
            }
        )?;

        linker.func_wrap(
            module,
            "message_get_data",
            |mut caller: wasmtime::Caller<'_, T>, id: u32, ptr: u32| {
                let ptr = ptr as usize;
                let len = caller.data().message_get_len(id) as usize;
                let dst = Self::get_memory_slice_bytes(&mut caller, ptr, len);
                caller.data().message_get_data(id, dst)
            }
        )?;

        Ok(())
    }

    fn get_panel_data(&self, panel: PanelId) -> u32 {
        *self.panel_datas.get(panel.0).unwrap()
    }

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

    fn get_memory_slice<D: bytemuck::Pod>(
        caller: &mut wasmtime::Caller<'_, T>,
        ptr: u32,
        num: u32,
    ) -> &'static mut [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_mut(bytes)
    }

    fn get_memory_slice_str(
        caller: &mut wasmtime::Caller<'_, T>,
        ptr: u32,
        len: u32,
    ) -> &'static mut str {
        let memory = Self::get_memory_slice_bytes(caller, ptr as usize, len as usize);
        std::str::from_utf8_mut(memory).unwrap()
    }

    fn get_memory_slice_bytes(
        caller: &mut wasmtime::Caller<'_, T>,
        ptr: usize,
        len: usize,
    ) -> &'static mut [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_mut(ptr, len)
        }
    }
}

impl<T: ScriptAbi> ScriptInstance for WasmtimeScript<T> {
    fn bind_panel(&mut self, msg: Vec<u8>) -> PanelId {
        let id = self.panel_datas.insert(0);
        let msg = self.store.data().message_new(msg);
        let args = (id as u32, msg);
        let data = self.bind_panel.call(&mut self.store, args).unwrap();
        *self.panel_datas.get_mut(id).unwrap() = data;
        self.store.data().message_free(msg);
        PanelId(id)
    }

    fn update(&mut self, panel: PanelId, dt: f32) {
        let data = self.get_panel_data(panel);
        self.update.call(&mut self.store, (data, dt)).unwrap();
    }

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

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

    fn on_message(&mut self, panel: PanelId, msg: Vec<u8>) {
        let data = self.get_panel_data(panel);
        let msg = self.store.data().message_new(msg);
        self.on_message.call(&mut self.store, (data, msg)).unwrap();
        self.store.data().message_free(msg);
    }
}

/// The standard [ScriptAbi] implementation to use.
#[derive(Default)]
pub struct ScriptAbiImpl {
    draw_cmds: Mutex<Vec<DrawCommand>>,
    font_store: text::FontStore,
    font_families: Mutex<HashMap<String, u32>>,
    loaded_fonts: RwLock<Vec<Arc<text::Font>>>,
    text_layouts: RwLock<Slab<text::TextLayout>>,
    message_store: RwLock<Slab<Vec<u8>>>,
}

impl ScriptAbi for ScriptAbiImpl {
    fn start_draw(&self) {
        let mut lock = self.draw_cmds.lock();
        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 draw_text_layout(&self, id: u32, offset: Vec2, scale: f32, color: Color) {
        // TODO multiple fonts per layout
        let layouts = self.text_layouts.read();
        let layout = layouts.get(id as usize).unwrap();
        let glyphs = layout.glyphs.as_slice();
        let loaded = self.loaded_fonts.read();
        let font = loaded.get(layout.font_id as usize).unwrap();
        let cmds = font.draw(glyphs, offset, scale, color);
        self.draw_cmds.lock().extend(cmds.into_iter());
    }

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

    fn font_load(&self, family: &str) -> u32 {
        let mut family_cache = self.font_families.lock();

        if let Some(cached) = family_cache.get(family) {
            return *cached;
        }

        let font = self.font_store.load_font(family);
        let mut loaded = self.loaded_fonts.write();
        let id = loaded.len() as u32;
        family_cache.insert(family.to_string(), id);
        loaded.push(font);
        id
    }

    fn text_layout_new(&self, font_id: u32, text: &str) -> u32 {
        let loaded = self.loaded_fonts.read();
        let font = loaded.get(font_id as usize).unwrap();
        let layout = font.shape(text);
        self.text_layouts.write().insert(layout) as u32
    }

    fn text_layout_delete(&self, id: u32) {
        self.text_layouts.write().remove(id as usize);
    }

    fn text_layout_get_bounds(&self, id: u32, dst: &mut Rect) {
        let src = self.text_layouts.read().get(id as usize).unwrap().bounds;
        let _ = std::mem::replace(dst, src);
    }

    fn message_new(&self, data: Vec<u8>) -> u32 {
        let mut store = self.message_store.write();
        let id = store.insert(data) as u32;
        id
    }

    fn message_free(&self, id: u32) {
        let mut store = self.message_store.write();
        store.remove(id as usize);
    }

    fn message_get_len(&self, id: u32) -> u32 {
        self.message_store.read().get(id as usize).unwrap().len() as u32
    }

    fn message_get_data(&self, id: u32, dst: &mut [u8]) {
        let store = self.message_store.read();
        let src = store.get(id as usize).unwrap();
        dst.copy_from_slice(src);
    }
}