// Copyright (c) 2022 Marceline Cramer
// SPDX-License-Identifier: LGPL-3.0-or-later

use std::collections::{hash_map::DefaultHasher, HashMap};
use std::hash::Hasher;
use std::ops::DerefMut;

use super::{Arc, Backend, Instance, PanelId, ScriptAbi};
use crate::DrawCommand;

use canary_script::{Color, CursorEventKind, Rect, Vec2};
use parking_lot::Mutex;
use prehash::{DefaultPrehasher, Prehashed, Prehasher};

type Caller<'a> = wasmtime::Caller<'a, Arc<ScriptAbi>>;
type Store = wasmtime::Store<Arc<ScriptAbi>>;
type Linker = wasmtime::Linker<Arc<ScriptAbi>>;
type ModuleCache = Mutex<HashMap<Prehashed<u64>, wasmtime::Module, DefaultPrehasher>>;

pub struct WasmtimeBackend {
    engine: wasmtime::Engine,
    module_cache: ModuleCache,
}

impl WasmtimeBackend {
    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);
        config.cache_config_load_default()?;

        let engine = wasmtime::Engine::new(&config)?;
        let module_cache = Default::default();
        Ok(Self {
            engine,
            module_cache,
        })
    }
}

impl Backend for WasmtimeBackend {
    fn load_module(&self, abi: Arc<ScriptAbi>, module: &[u8]) -> anyhow::Result<Arc<dyn Instance>> {
        let mut hasher = DefaultHasher::new();
        hasher.write(module);
        let hashed = hasher.finish();

        let prehasher = DefaultPrehasher::new();
        let prehashed = prehasher.prehash(hashed);
        let mut cache = self.module_cache.lock();

        let module = if let Some(module) = cache.get(&prehashed) {
            module
        } else {
            let module = wasmtime::Module::new(&self.engine, module)?;
            cache.insert(prehashed, module);
            cache.get(&prehashed).unwrap()
        };

        let mut store = wasmtime::Store::new(&self.engine, abi);
        let mut linker = Linker::new(&self.engine);
        WasmtimeInstance::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_resize = instance.get_typed_func(&mut store, "on_resize")?;
        let on_cursor_event = instance.get_typed_func(&mut store, "on_cursor_event")?;
        let on_message = instance.get_typed_func(&mut store, "on_message")?;

        let instance = WasmtimeInstance {
            store: Mutex::new(store),
            bind_panel,
            update,
            draw,
            on_resize,
            on_cursor_event,
            on_message,
        };

        let instance = Arc::new(instance);

        Ok(instance)
    }
}

pub struct WasmtimeInstance {
    store: Mutex<Store>,
    bind_panel: wasmtime::TypedFunc<(u32, u32, u32), u32>,
    update: wasmtime::TypedFunc<(u32, f32), ()>,
    draw: wasmtime::TypedFunc<u32, ()>,
    on_resize: wasmtime::TypedFunc<(u32, f32, f32), ()>,
    on_cursor_event: wasmtime::TypedFunc<(u32, u32, f32, f32), ()>,
    on_message: wasmtime::TypedFunc<(u32, u32), ()>,
}

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

        linker.func_wrap(
            module,
            "draw_indexed",
            |mut caller: Caller<'_>,
             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: Caller<'_>, id: u32, xoff: f32, yoff: f32, scale: f32, color: u32| {
                let offset = Vec2 { x: xoff, y: yoff };
                let color = Color(color);
                caller.data().draw_text_layout(id, offset, scale, color);
            },
        )?;

        linker.func_wrap(
            module,
            "font_load",
            |mut caller: Caller<'_>, 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: Caller<'_>, 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: Caller<'_>, id: u32| caller.data().text_layout_delete(id),
        )?;

        linker.func_wrap(
            module,
            "text_layout_get_bounds",
            |mut caller: Caller<'_>, 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: Caller<'_>, id: u32| -> u32 { caller.data().message_get_len(id) },
        )?;

        linker.func_wrap(
            module,
            "message_get_data",
            |mut caller: Caller<'_>, 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)
            },
        )?;

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

        Ok(())
    }

    fn get_memory_ref<D: bytemuck::Pod>(caller: &mut Caller<'_>, 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 Caller<'_>,
        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 Caller<'_>, 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 Caller<'_>,
        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 Instance for WasmtimeInstance {
    fn bind_panel(&self, panel: PanelId, protocol: &str, msg: Vec<u8>) -> u32 {
        let mut store = self.store.lock();
        let protocol = store.data().message_new(protocol.as_bytes().to_vec());
        let msg = store.data().message_new(msg);
        let args = (panel.0 as u32, protocol, msg);
        let data = self.bind_panel.call(store.deref_mut(), args).unwrap();
        store.data().message_free(protocol);
        store.data().message_free(msg);
        data
    }

    fn update(&self, panel_ud: u32, dt: f32) {
        let mut store = self.store.lock();
        self.update.call(store.deref_mut(), (panel_ud, dt)).unwrap();
    }

    fn draw(&self, panel_ud: u32) -> Vec<DrawCommand> {
        let mut store = self.store.lock();
        store.data().start_draw();
        self.draw.call(store.deref_mut(), panel_ud).unwrap();
        let mut cmds = Vec::new();
        store
            .data()
            .with_draw_commands(|slice| cmds.extend_from_slice(slice));
        cmds
    }

    fn on_resize(&self, panel_ud: u32, new_size: Vec2) {
        let mut store = self.store.lock();
        self.on_resize
            .call(store.deref_mut(), (panel_ud, new_size.x, new_size.y))
            .unwrap();
    }

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

    fn on_message(&self, panel_ud: u32, msg: Vec<u8>) {
        let mut store = self.store.lock();
        let msg = store.data().message_new(msg);
        self.on_message
            .call(store.deref_mut(), (panel_ud, msg))
            .unwrap();
        store.data().message_free(msg);
    }
}