use anyhow::Result;
use nightshade::editor_plugin::{
    EditorEvent, InspectorRequest, InspectorResponse, MenuRequest, MenuResponse, ModCommand,
    ModInfo, UiElement, egui,
};
use std::collections::HashMap;
use std::path::Path;
use wasmtime::*;

struct PluginInstance {
    id: u64,
    store: Store<PluginState>,
    instance: Instance,
    on_init: Option<TypedFunc<(), ()>>,
    on_frame: Option<TypedFunc<(), ()>>,
    info: Option<ModInfo>,
}

struct PluginState {
    wasi: wasmtime_wasi::preview1::WasiP1Ctx,
    pending_commands: Vec<ModCommand>,
    pending_inspector_response: Option<InspectorResponse>,
    pending_menu_response: Option<MenuResponse>,
    pending_mod_info: Option<ModInfo>,
}

pub struct RegisteredInspector {
    pub name: String,
    pub plugin_id: u64,
}

pub struct RegisteredMenuItem {
    pub path: String,
    pub plugin_id: u64,
}

pub struct EditorPluginRuntime {
    engine: Engine,
    plugins: Vec<PluginInstance>,
    next_plugin_id: u64,
    registered_inspectors: HashMap<String, RegisteredInspector>,
    registered_menu_items: HashMap<String, RegisteredMenuItem>,
}

impl EditorPluginRuntime {
    pub fn new() -> Result<Self> {
        Ok(Self {
            engine: Engine::default(),
            plugins: Vec::new(),
            next_plugin_id: 1,
            registered_inspectors: HashMap::new(),
            registered_menu_items: HashMap::new(),
        })
    }

    pub fn load_plugins_from_directory(&mut self, dir: &Path) -> Result<()> {
        if !dir.exists() {
            tracing::info!("Plugin directory does not exist: {:?}", dir);
            return Ok(());
        }

        for entry in std::fs::read_dir(dir)? {
            let entry = entry?;
            let path = entry.path();
            if path.extension().is_some_and(|ext| ext == "wasm") {
                match self.load_plugin(&path) {
                    Ok(_) => tracing::info!("Loaded editor plugin: {:?}", path),
                    Err(error) => {
                        tracing::error!("Failed to load editor plugin {:?}: {}", path, error)
                    }
                }
            }
        }

        Ok(())
    }

    fn load_plugin(&mut self, path: &Path) -> Result<()> {
        let module = Module::from_file(&self.engine, path)?;
        let mut linker: Linker<PluginState> = Linker::new(&self.engine);

        wasmtime_wasi::preview1::add_to_linker_sync(&mut linker, |state| &mut state.wasi)?;

        linker.func_wrap(
            "env",
            "host_send_mod_command",
            |mut caller: Caller<'_, PluginState>, ptr: u32, len: u32| {
                let Some(bytes) = read_plugin_memory(&mut caller, ptr, len) else {
                    return;
                };
                if let Some(cmd) = ModCommand::from_bytes(&bytes) {
                    caller.data_mut().pending_commands.push(cmd);
                }
            },
        )?;

        linker.func_wrap(
            "env",
            "host_send_mod_info",
            |mut caller: Caller<'_, PluginState>, ptr: u32, len: u32| {
                let Some(bytes) = read_plugin_memory(&mut caller, ptr, len) else {
                    return;
                };
                if let Some(info) = ModInfo::from_bytes(&bytes) {
                    caller.data_mut().pending_mod_info = Some(info);
                }
            },
        )?;

        linker.func_wrap(
            "env",
            "host_send_inspector_response",
            |mut caller: Caller<'_, PluginState>, ptr: u32, len: u32| {
                let Some(bytes) = read_plugin_memory(&mut caller, ptr, len) else {
                    return;
                };
                if let Some(response) = InspectorResponse::from_bytes(&bytes) {
                    caller.data_mut().pending_inspector_response = Some(response);
                }
            },
        )?;

        linker.func_wrap(
            "env",
            "host_send_menu_response",
            |mut caller: Caller<'_, PluginState>, ptr: u32, len: u32| {
                let Some(bytes) = read_plugin_memory(&mut caller, ptr, len) else {
                    return;
                };
                if let Some(response) = MenuResponse::from_bytes(&bytes) {
                    caller.data_mut().pending_menu_response = Some(response);
                }
            },
        )?;

        let wasi = wasmtime_wasi::WasiCtxBuilder::new()
            .inherit_stdio()
            .build_p1();

        let state = PluginState {
            wasi,
            pending_commands: Vec::new(),
            pending_inspector_response: None,
            pending_menu_response: None,
            pending_mod_info: None,
        };

        let mut store = Store::new(&self.engine, state);
        let instance = linker.instantiate(&mut store, &module)?;

        let on_init = instance
            .get_typed_func::<(), ()>(&mut store, "on_init")
            .ok();
        let on_frame = instance
            .get_typed_func::<(), ()>(&mut store, "on_frame")
            .ok();

        let plugin_id = self.next_plugin_id;
        self.next_plugin_id += 1;

        self.plugins.push(PluginInstance {
            id: plugin_id,
            store,
            instance,
            on_init,
            on_frame,
            info: None,
        });

        Ok(())
    }

    pub fn call_on_init(&mut self) {
        for plugin in &mut self.plugins {
            if let Some(ref on_init) = plugin.on_init {
                if let Err(error) = on_init.call(&mut plugin.store, ()) {
                    tracing::error!("Plugin {} on_init failed: {}", plugin.id, error);
                }
                if let Some(info) = plugin.store.data_mut().pending_mod_info.take() {
                    tracing::info!("Plugin registered: {} v{}", info.name, info.version);
                    plugin.info = Some(info);
                }
            }
        }
        self.process_pending_commands();
    }

    pub fn call_on_frame(&mut self) {
        for plugin in &mut self.plugins {
            if let Some(ref on_frame) = plugin.on_frame
                && let Err(error) = on_frame.call(&mut plugin.store, ())
            {
                tracing::error!("Plugin {} on_frame failed: {}", plugin.id, error);
            }
        }
        self.process_pending_commands();
    }

    fn process_pending_commands(&mut self) {
        for plugin in &mut self.plugins {
            for command in plugin.store.data_mut().pending_commands.drain(..) {
                match command {
                    ModCommand::Log { message } => {
                        tracing::info!("[Plugin {}] {}", plugin.id, message);
                    }
                    ModCommand::RegisterInspector { name, inspector_id } => {
                        tracing::info!(
                            "Plugin {} registered inspector '{}' ({})",
                            plugin.id,
                            name,
                            inspector_id
                        );
                        self.registered_inspectors.insert(
                            inspector_id,
                            RegisteredInspector {
                                name,
                                plugin_id: plugin.id,
                            },
                        );
                    }
                    ModCommand::RegisterMenuItem { path, item_id } => {
                        tracing::info!(
                            "Plugin {} registered menu item '{}' ({})",
                            plugin.id,
                            path,
                            item_id
                        );
                        self.registered_menu_items.insert(
                            item_id,
                            RegisteredMenuItem {
                                path,
                                plugin_id: plugin.id,
                            },
                        );
                    }
                }
            }
        }
    }

    pub fn send_event(&mut self, event: &EditorEvent) {
        let Ok(bytes) = event.to_bytes() else {
            return;
        };

        for plugin in &mut self.plugins {
            let alloc_fn = plugin
                .instance
                .get_typed_func::<u32, u32>(&mut plugin.store, "mod_alloc");
            let receive_fn = plugin
                .instance
                .get_typed_func::<(u32, u32), ()>(&mut plugin.store, "mod_receive_event");

            let (Ok(alloc_fn), Ok(receive_fn)) = (alloc_fn, receive_fn) else {
                continue;
            };

            let Ok(ptr) = alloc_fn.call(&mut plugin.store, bytes.len() as u32) else {
                continue;
            };

            let Some(memory) = plugin.instance.get_memory(&mut plugin.store, "memory") else {
                continue;
            };

            let data = memory.data_mut(&mut plugin.store);
            if (ptr as usize) + bytes.len() <= data.len() {
                data[ptr as usize..ptr as usize + bytes.len()].copy_from_slice(&bytes);
                let _ = receive_fn.call(&mut plugin.store, (ptr, bytes.len() as u32));
            }
        }
    }

    pub fn request_inspector_ui(
        &mut self,
        inspector_id: &str,
        entity_id: u64,
    ) -> Option<Vec<UiElement>> {
        let registered = self.registered_inspectors.get(inspector_id)?;
        let plugin_id = registered.plugin_id;
        let plugin = self.plugins.iter_mut().find(|p| p.id == plugin_id)?;

        let request = InspectorRequest::RenderUi {
            inspector_id: inspector_id.to_string(),
            entity_id,
        };
        let Ok(bytes) = request.to_bytes() else {
            return None;
        };

        let alloc_fn = plugin
            .instance
            .get_typed_func::<u32, u32>(&mut plugin.store, "mod_request_alloc")
            .ok()?;
        let receive_fn = plugin
            .instance
            .get_typed_func::<(u32, u32), ()>(&mut plugin.store, "mod_receive_inspector_request")
            .ok()?;

        let ptr = alloc_fn.call(&mut plugin.store, bytes.len() as u32).ok()?;
        let memory = plugin.instance.get_memory(&mut plugin.store, "memory")?;

        let data = memory.data_mut(&mut plugin.store);
        if (ptr as usize) + bytes.len() <= data.len() {
            data[ptr as usize..ptr as usize + bytes.len()].copy_from_slice(&bytes);
            receive_fn
                .call(&mut plugin.store, (ptr, bytes.len() as u32))
                .ok()?;
        }

        if let Some(ref on_frame) = plugin.on_frame {
            let _ = on_frame.call(&mut plugin.store, ());
        }

        let response = plugin.store.data_mut().pending_inspector_response.take()?;
        match response {
            InspectorResponse::Ui { elements } => Some(elements),
            InspectorResponse::HasComponent { .. } => None,
        }
    }

    pub fn get_registered_inspectors(&self) -> Vec<(String, String)> {
        self.registered_inspectors
            .iter()
            .map(|(id, registered)| (id.clone(), registered.name.clone()))
            .collect()
    }

    pub fn get_registered_menu_items(&self) -> Vec<(String, String)> {
        self.registered_menu_items
            .iter()
            .map(|(id, registered)| (id.clone(), registered.path.clone()))
            .collect()
    }

    pub fn request_menu_ui(&mut self, item_id: &str) -> Option<Vec<UiElement>> {
        let registered = self.registered_menu_items.get(item_id)?;
        let plugin_id = registered.plugin_id;
        let plugin = self.plugins.iter_mut().find(|p| p.id == plugin_id)?;

        let request = MenuRequest::RenderUi {
            item_id: item_id.to_string(),
        };
        let Ok(bytes) = request.to_bytes() else {
            return None;
        };

        let alloc_fn = plugin
            .instance
            .get_typed_func::<u32, u32>(&mut plugin.store, "mod_request_alloc")
            .ok()?;
        let receive_fn = plugin
            .instance
            .get_typed_func::<(u32, u32), ()>(&mut plugin.store, "mod_receive_menu_request")
            .ok()?;

        let ptr = alloc_fn.call(&mut plugin.store, bytes.len() as u32).ok()?;
        let memory = plugin.instance.get_memory(&mut plugin.store, "memory")?;

        let data = memory.data_mut(&mut plugin.store);
        if (ptr as usize) + bytes.len() <= data.len() {
            data[ptr as usize..ptr as usize + bytes.len()].copy_from_slice(&bytes);
            receive_fn
                .call(&mut plugin.store, (ptr, bytes.len() as u32))
                .ok()?;
        }

        if let Some(ref on_frame) = plugin.on_frame {
            let _ = on_frame.call(&mut plugin.store, ());
        }

        let response = plugin.store.data_mut().pending_menu_response.take()?;
        match response {
            MenuResponse::Ui { elements } => Some(elements),
        }
    }

    pub fn get_loaded_plugins(&self) -> Vec<ModInfo> {
        self.plugins.iter().filter_map(|p| p.info.clone()).collect()
    }

    #[allow(dead_code)]
    pub fn plugin_count(&self) -> usize {
        self.plugins.len()
    }
}

fn read_plugin_memory(caller: &mut Caller<'_, PluginState>, ptr: u32, len: u32) -> Option<Vec<u8>> {
    let memory = caller.get_export("memory").and_then(|e| e.into_memory())?;
    let data = memory.data(caller);
    if (ptr as usize) + (len as usize) > data.len() {
        return None;
    }
    Some(data[ptr as usize..(ptr + len) as usize].to_vec())
}

pub fn render_ui_elements(ui: &mut egui::Ui, elements: &[UiElement]) {
    for element in elements {
        render_ui_element(ui, element);
    }
}

fn render_ui_element(ui: &mut egui::Ui, element: &UiElement) {
    match element {
        UiElement::Label { text } => {
            ui.label(text);
        }
        UiElement::Button { text, id: _ } => {
            let _ = ui.button(text);
        }
        UiElement::TextEdit { value, id: _ } => {
            let mut text = value.clone();
            ui.text_edit_singleline(&mut text);
        }
        UiElement::Separator => {
            ui.separator();
        }
        UiElement::Spacing { amount } => {
            ui.add_space(*amount);
        }
        UiElement::Horizontal { children } => {
            ui.horizontal(|ui| {
                render_ui_elements(ui, children);
            });
        }
        UiElement::Vertical { children } => {
            ui.vertical(|ui| {
                render_ui_elements(ui, children);
            });
        }
        UiElement::Group { children } => {
            ui.group(|ui| {
                render_ui_elements(ui, children);
            });
        }
        UiElement::Collapsing {
            header,
            default_open,
            children,
        } => {
            egui::CollapsingHeader::new(header)
                .default_open(*default_open)
                .show(ui, |ui| {
                    render_ui_elements(ui, children);
                });
        }
    }
}
