feat(plugins): update and render plugins asynchronously (#2410)

* working-ish minus a few race conditions

* relax atomicity

* various refactoringz

* remove commented out code

* clarify some stuffs

* refactor(plugins): move PluginMap and friends to a different file

* refactor(plugins): move zellij_exports and friends to a different file

* style(fmt): rustfmt

* fix(e2e): adjust tests for flakiness async loading

5 files changed, 791 insertions, 565 deletions

diff --git a/zellij-server/src/plugins/mod.rs b/zellij-server/src/plugins/mod.rs
index eeb19c20..aaf6164b 100644
--- a/zellij-server/src/plugins/mod.rs
+++ b/zellij-server/src/plugins/mod.rs
@@ -1,5 +1,7 @@
 mod plugin_loader;
+mod plugin_map;
 mod wasm_bridge;
+mod zellij_exports;
 use log::info;
 use std::{collections::HashMap, fs, path::PathBuf};
 use wasmer::Store;
@@ -37,7 +39,6 @@ pub enum PluginInstruction {
         Option<String>, // pane title
         RunPlugin,
         usize, // tab index
-        ClientId,
         Size,
     ),
     Resize(u32, usize, usize), // plugin_id, columns, rows
@@ -91,7 +92,7 @@ pub(crate) fn plugin_thread_main(
         err_ctx.add_call(ContextType::Plugin((&event).into()));
         match event {
             PluginInstruction::Load(should_float, pane_title, run, tab_index, client_id, size) => {
-                match wasm_bridge.load_plugin(&run, tab_index, size, client_id) {
+                match wasm_bridge.load_plugin(&run, tab_index, size, Some(client_id)) {
                     Ok(plugin_id) => {
                         drop(bus.senders.send_to_screen(ScreenInstruction::AddPlugin(
                             should_float,
@@ -112,41 +113,38 @@ pub(crate) fn plugin_thread_main(
             PluginInstruction::Unload(pid) => {
                 wasm_bridge.unload_plugin(pid)?;
             },
-            PluginInstruction::Reload(
-                should_float,
-                pane_title,
-                run,
-                tab_index,
-                client_id,
-                size,
-            ) => match wasm_bridge.reload_plugin(&run) {
-                Ok(_) => {
-                    let _ = bus
-                        .senders
-                        .send_to_server(ServerInstruction::UnblockInputThread);
-                },
-                Err(err) => match err.downcast_ref::<ZellijError>() {
-                    Some(ZellijError::PluginDoesNotExist) => {
-                        log::warn!("Plugin {} not found, starting it instead", run.location);
-                        match wasm_bridge.load_plugin(&run, tab_index, size, client_id) {
-                            Ok(plugin_id) => {
-                                drop(bus.senders.send_to_screen(ScreenInstruction::AddPlugin(
-                                    should_float,
-                                    run,
-                                    pane_title,
-                                    tab_index,
-                                    plugin_id,
-                                )));
-                            },
-                            Err(e) => {
-                                log::error!("Failed to load plugin: {e}");
-                            },
-                        };
+            PluginInstruction::Reload(should_float, pane_title, run, tab_index, size) => {
+                match wasm_bridge.reload_plugin(&run) {
+                    Ok(_) => {
+                        let _ = bus
+                            .senders
+                            .send_to_server(ServerInstruction::UnblockInputThread);
                     },
-                    _ => {
-                        return Err(err);
+                    Err(err) => match err.downcast_ref::<ZellijError>() {
+                        Some(ZellijError::PluginDoesNotExist) => {
+                            log::warn!("Plugin {} not found, starting it instead", run.location);
+                            // we intentionally do not provide the client_id here because it belongs to
+                            // the cli who spawned the command and is not an existing client_id
+                            match wasm_bridge.load_plugin(&run, tab_index, size, None) {
+                                Ok(plugin_id) => {
+                                    drop(bus.senders.send_to_screen(ScreenInstruction::AddPlugin(
+                                        should_float,
+                                        run,
+                                        pane_title,
+                                        tab_index,
+                                        plugin_id,
+                                    )));
+                                },
+                                Err(e) => {
+                                    log::error!("Failed to load plugin: {e}");
+                                },
+                            };
+                        },
+                        _ => {
+                            return Err(err);
+                        },
                     },
-                },
+                }
             },
             PluginInstruction::Resize(pid, new_columns, new_rows) => {
                 wasm_bridge.resize_plugin(pid, new_columns, new_rows)?;
@@ -184,7 +182,7 @@ pub(crate) fn plugin_thread_main(
                 for run_instruction in extracted_run_instructions {
                     if let Some(Run::Plugin(run)) = run_instruction {
                         let plugin_id =
-                            wasm_bridge.load_plugin(&run, tab_index, size, client_id)?;
+                            wasm_bridge.load_plugin(&run, tab_index, size, Some(client_id))?;
                         plugin_ids.entry(run.location).or_default().push(plugin_id);
                     }
                 }
diff --git a/zellij-server/src/plugins/plugin_loader.rs b/zellij-server/src/plugins/plugin_loader.rs
index c5413388..86afd8ca 100644
--- a/zellij-server/src/plugins/plugin_loader.rs
+++ b/zellij-server/src/plugins/plugin_loader.rs
@@ -1,4 +1,5 @@
-use crate::plugins::wasm_bridge::{wasi_read_string, zellij_exports, PluginEnv, PluginMap};
+use crate::plugins::plugin_map::{PluginEnv, PluginMap, RunningPlugin, Subscriptions};
+use crate::plugins::zellij_exports::{wasi_read_string, zellij_exports};
 use highway::{HighwayHash, PortableHash};
 use log::info;
 use semver::Version;
@@ -146,21 +147,8 @@ fn assert_plugin_version(instance: &Instance, plugin_env: &PluginEnv) -> Result<
     Ok(())
 }
 
-fn load_plugin_instance(instance: &mut Instance) -> Result<()> {
-    let err_context = || format!("failed to load plugin from instance {instance:#?}");
-
-    let load_function = instance
-        .exports
-        .get_function("_start")
-        .with_context(err_context)?;
-    // This eventually calls the `.load()` method
-    load_function.call(&[]).with_context(err_context)?;
-    Ok(())
-}
-
 pub struct PluginLoader<'a> {
     plugin_cache: Arc<Mutex<HashMap<PathBuf, Module>>>,
-    plugin_map: Arc<Mutex<PluginMap>>,
     plugin_path: PathBuf,
     loading_indication: &'a mut LoadingIndication,
     senders: ThreadSenders,
@@ -206,15 +194,20 @@ impl<'a> PluginLoader<'a> {
         )?;
         plugin_loader
             .load_module_from_memory()
-            .and_then(|module| plugin_loader.create_plugin_instance_and_environment(module))
-            .and_then(|(instance, plugin_env)| {
-                plugin_loader.load_plugin_instance(&instance, &plugin_env)?;
-                plugin_loader.clone_instance_for_other_clients(
+            .and_then(|module| {
+                plugin_loader.create_plugin_instance_environment_and_subscriptions(module)
+            })
+            .and_then(|(instance, plugin_env, subscriptions)| {
+                plugin_loader.load_plugin_instance(
                     &instance,
                     &plugin_env,
-                    &connected_clients,
+                    &plugin_map,
+                    &subscriptions,
                 )
             })
+            .and_then(|_| {
+                plugin_loader.clone_instance_for_other_clients(&connected_clients, &plugin_map)
+            })
             .with_context(err_context)?;
         display_loading_stage!(end, loading_indication, senders, plugin_id);
         Ok(())
@@ -237,7 +230,6 @@ impl<'a> PluginLoader<'a> {
         let err_context = || format!("failed to start plugin {plugin:#?} for client {client_id}");
         let mut plugin_loader = PluginLoader::new(
             &plugin_cache,
-            &plugin_map,
             loading_indication,
             &senders,
             plugin_id,
@@ -252,19 +244,69 @@ impl<'a> PluginLoader<'a> {
             .load_module_from_memory()
             .or_else(|_e| plugin_loader.load_module_from_hd_cache())
             .or_else(|_e| plugin_loader.compile_module())
-            .and_then(|module| plugin_loader.create_plugin_instance_and_environment(module))
-            .and_then(|(instance, plugin_env)| {
-                plugin_loader.load_plugin_instance(&instance, &plugin_env)?;
-                plugin_loader.clone_instance_for_other_clients(
+            .and_then(|module| {
+                plugin_loader.create_plugin_instance_environment_and_subscriptions(module)
+            })
+            .and_then(|(instance, plugin_env, subscriptions)| {
+                plugin_loader.load_plugin_instance(
                     &instance,
                     &plugin_env,
+                    &plugin_map,
+                    &subscriptions,
+                )
+            })
+            .and_then(|_| {
+                plugin_loader.clone_instance_for_other_clients(
                     &connected_clients.lock().unwrap(),
+                    &plugin_map,
                 )
             })
             .with_context(err_context)?;
         display_loading_stage!(end, loading_indication, senders, plugin_id);
         Ok(())
     }
+    pub fn add_client(
+        client_id: ClientId,
+        plugin_dir: PathBuf,
+        plugin_cache: Arc<Mutex<HashMap<PathBuf, Module>>>,
+        senders: ThreadSenders,
+        store: Store,
+        plugin_map: Arc<Mutex<PluginMap>>,
+        connected_clients: Arc<Mutex<Vec<ClientId>>>,
+        loading_indication: &mut LoadingIndication,
+    ) -> Result<()> {
+        let mut new_plugins = HashSet::new();
+        for (&(plugin_id, _), _) in &*plugin_map.lock().unwrap() {
+            new_plugins.insert((plugin_id, client_id));
+        }
+        for (plugin_id, existing_client_id) in new_plugins {
+            let mut plugin_loader = PluginLoader::new_from_different_client_id(
+                &plugin_cache,
+                &plugin_map,
+                loading_indication,
+                &senders,
+                plugin_id,
+                existing_client_id,
+                &store,
+                &plugin_dir,
+            )?;
+            plugin_loader
+                .load_module_from_memory()
+                .and_then(|module| {
+                    plugin_loader.create_plugin_instance_environment_and_subscriptions(module)
+                })
+                .and_then(|(instance, plugin_env, subscriptions)| {
+                    plugin_loader.load_plugin_instance(
+                        &instance,
+                        &plugin_env,
+                        &plugin_map,
+                        &subscriptions,
+                    )
+                })?
+        }
+        connected_clients.lock().unwrap().push(client_id);
+        Ok(())
+    }
 
     pub fn reload_plugin(
         plugin_id: u32,
@@ -297,22 +339,26 @@ impl<'a> PluginLoader<'a> {
         )?;
         plugin_loader
             .compile_module()
-            .and_then(|module| plugin_loader.create_plugin_instance_and_environment(module))
-            .and_then(|(instance, plugin_env)| {
-                plugin_loader.load_plugin_instance(&instance, &plugin_env)?;
-                plugin_loader.clone_instance_for_other_clients(
+            .and_then(|module| {
+                plugin_loader.create_plugin_instance_environment_and_subscriptions(module)
+            })
+            .and_then(|(instance, plugin_env, subscriptions)| {
+                plugin_loader.load_plugin_instance(
                     &instance,
                     &plugin_env,
-                    &connected_clients,
+                    &plugin_map,
+                    &subscriptions,
                 )
             })
+            .and_then(|_| {
+                plugin_loader.clone_instance_for_other_clients(&connected_clients, &plugin_map)
+            })
             .with_context(err_context)?;
         display_loading_stage!(end, loading_indication, senders, plugin_id);
         Ok(())
     }
     pub fn new(
         plugin_cache: &Arc<Mutex<HashMap<PathBuf, Module>>>,
-        plugin_map: &Arc<Mutex<PluginMap>>,
         loading_indication: &'a mut LoadingIndication,
         senders: &ThreadSenders,
         plugin_id: u32,
@@ -328,7 +374,6 @@ impl<'a> PluginLoader<'a> {
         let plugin_path = plugin.path.clone();
         Ok(PluginLoader {
             plugin_cache: plugin_cache.clone(),
-            plugin_map: plugin_map.clone(),
             plugin_path,
             loading_indication,
             senders: senders.clone(),
@@ -354,19 +399,63 @@ impl<'a> PluginLoader<'a> {
         plugin_dir: &'a PathBuf,
     ) -> Result<Self> {
         let err_context = || "Failed to find existing plugin";
-        let (_old_instance, old_user_env, (rows, cols)) = {
+        let (running_plugin, _subscriptions) = {
             let mut plugin_map = plugin_map.lock().unwrap();
             plugin_map
                 .remove(&(plugin_id, client_id))
                 .with_context(err_context)?
         };
-        let tab_index = old_user_env.tab_index;
-        let size = Size { rows, cols };
-        let plugin_config = old_user_env.plugin.clone();
-        loading_indication.set_name(old_user_env.name());
+        let running_plugin = running_plugin.lock().unwrap();
+        let tab_index = running_plugin.plugin_env.tab_index;
+        let size = Size {
+            rows: running_plugin.rows,
+            cols: running_plugin.columns,
+        };
+        let plugin_config = running_plugin.plugin_env.plugin.clone();
+        loading_indication.set_name(running_plugin.plugin_env.name());
+        PluginLoader::new(
+            plugin_cache,
+            loading_indication,
+            senders,
+            plugin_id,
+            client_id,
+            store,
+            plugin_config,
+            plugin_dir,
+            tab_index,
+            size,
+        )
+    }
+    pub fn new_from_different_client_id(
+        plugin_cache: &Arc<Mutex<HashMap<PathBuf, Module>>>,
+        plugin_map: &Arc<Mutex<PluginMap>>,
+        loading_indication: &'a mut LoadingIndication,
+        senders: &ThreadSenders,
+        plugin_id: u32,
+        client_id: ClientId,
+        store: &Store,
+        plugin_dir: &'a PathBuf,
+    ) -> Result<Self> {
+        let err_context = || "Failed to find existing plugin";
+        let (running_plugin, _subscriptions) = {
+            let plugin_map = plugin_map.lock().unwrap();
+            plugin_map
+                .iter()
+                .find(|((p_id, _c_id), _)| p_id == &plugin_id)
+                .with_context(err_context)?
+                .1
+                .clone()
+        };
+        let running_plugin = running_plugin.lock().unwrap();
+        let tab_index = running_plugin.plugin_env.tab_index;
+        let size = Size {
+            rows: running_plugin.rows,
+            cols: running_plugin.columns,
+        };
+        let plugin_config = running_plugin.plugin_env.plugin.clone();
+        loading_indication.set_name(running_plugin.plugin_env.name());
         PluginLoader::new(
             plugin_cache,
-            plugin_map,
             loading_indication,
             senders,
             plugin_id,
@@ -464,10 +553,10 @@ impl<'a> PluginLoader<'a> {
             .with_context(err_context)?;
         Ok(module)
     }
-    pub fn create_plugin_instance_and_environment(
+    pub fn create_plugin_instance_environment_and_subscriptions(
         &mut self,
         module: Module,
-    ) -> Result<(Instance, PluginEnv)> {
+    ) -> Result<(Instance, PluginEnv, Arc<Mutex<Subscriptions>>)> {
         let err_context = || {
             format!(
                 "Failed to create instance and plugin env for plugin {}",
@@ -499,12 +588,12 @@ impl<'a> PluginLoader<'a> {
             plugin: mut_plugin,
             senders: self.senders.clone(),
             wasi_env,
-            subscriptions: Arc::new(Mutex::new(HashSet::new())),
             plugin_own_data_dir: self.plugin_own_data_dir.clone(),
             tab_index: self.tab_index,
         };
 
-        let zellij = zellij_exports(&self.store, &plugin_env);
+        let subscriptions = Arc::new(Mutex::new(HashSet::new()));
+        let zellij = zellij_exports(&self.store, &plugin_env, &subscriptions);
         let instance =
             Instance::new(&module, &zellij.chain_back(wasi)).with_context(err_context)?;
         assert_plugin_version(&instance, &plugin_env).with_context(err_context)?;
@@ -514,12 +603,14 @@ impl<'a> PluginLoader<'a> {
             .lock()
             .unwrap()
             .insert(cloned_plugin.path, module);
-        Ok((instance, plugin_env))
+        Ok((instance, plugin_env, subscriptions))
     }
     pub fn load_plugin_instance(
         &mut self,
         instance: &Instance,
         plugin_env: &PluginEnv,
+        plugin_map: &Arc<Mutex<PluginMap>>,
+        subscriptions: &Arc<Mutex<Subscriptions>>,
     ) -> Result<()> {
         let err_context = || format!("failed to load plugin from instance {instance:#?}");
         let main_user_instance = instance.clone();
@@ -548,13 +639,16 @@ impl<'a> PluginLoader<'a> {
             self.senders,
             self.plugin_id
         );
-        let mut plugin_map = self.plugin_map.lock().unwrap();
-        plugin_map.insert(
+        plugin_map.lock().unwrap().insert(
             (self.plugin_id, self.client_id),
             (
-                main_user_instance,
-                main_user_env,
-                (self.size.rows, self.size.cols),
+                Arc::new(Mutex::new(RunningPlugin::new(
+                    main_user_instance,
+                    main_user_env,
+                    self.size.rows,
+                    self.size.cols,
+                ))),
+                subscriptions.clone(),
             ),
         );
         display_loading_stage!(
@@ -567,9 +661,8 @@ impl<'a> PluginLoader<'a> {
     }
     pub fn clone_instance_for_other_clients(
         &mut self,
-        instance: &Instance,
-        plugin_env: &PluginEnv,
         connected_clients: &[ClientId],
+        plugin_map: &Arc<Mutex<PluginMap>>,
     ) -> Result<()> {
         if !connected_clients.is_empty() {
             display_loading_stage!(
@@ -578,14 +671,32 @@ impl<'a> PluginLoader<'a> {
                 self.senders,
                 self.plugin_id
             );
-            let mut plugin_map = self.plugin_map.lock().unwrap();
             for client_id in connected_clients {
-                let (instance, new_plugin_env) =
-                    clone_plugin_for_client(&plugin_env, *client_id, &instance, &self.store)?;
-                plugin_map.insert(
-                    (self.plugin_id, *client_id),
-                    (instance, new_plugin_env, (self.size.rows, self.size.cols)),
-                );
+                let mut loading_indication = LoadingIndication::new("".into());
+                let mut plugin_loader_for_client = PluginLoader::new_from_different_client_id(
+                    &self.plugin_cache.clone(),
+                    &plugin_map,
+                    &mut loading_indication,
+                    &self.senders.clone(),
+                    self.plugin_id,
+                    *client_id,
+                    &self.store,
+                    &self.plugin_dir,
+                )?;
+                plugin_loader_for_client
+                    .load_module_from_memory()
+                    .and_then(|module| {
+                        plugin_loader_for_client
+                            .create_plugin_instance_environment_and_subscriptions(module)
+                    })
+                    .and_then(|(instance, plugin_env, subscriptions)| {
+                        plugin_loader_for_client.load_plugin_instance(
+                            &instance,
+                            &plugin_env,
+                            plugin_map,
+                            &subscriptions,
+                        )
+                    })?
             }
             display_loading_stage!(
                 indicate_cloning_plugin_for_other_clients_success,
@@ -633,24 +744,3 @@ fn create_plugin_fs_entries(plugin_own_data_dir: &PathBuf) -> Result<()> {
         .with_context(err_context)?;
     Ok(())
 }
-
-fn clone_plugin_for_client(
-    plugin_env: &PluginEnv,
-    client_id: ClientId,
-    instance: &Instance,
-    store: &Store,
-) -> Result<(Instance, PluginEnv)> {
-    let err_context = || format!("Failed to clone plugin for client {client_id}");
-    let mut new_plugin_env = plugin_env.clone();
-    new_plugin_env.client_id = client_id;
-    let module = instance.module().clone();
-    let wasi = new_plugin_env
-        .wasi_env
-        .import_object(&module)
-        .with_context(err_context)?;
-    let zellij = zellij_exports(store, &new_plugin_env);
-    let mut instance =
-        Instance::new(&module, &zellij.chain_back(wasi)).with_context(err_context)?;
-    load_plugin_instance(&mut instance).with_context(err_context)?;
-    Ok((instance, new_plugin_env))
-}
diff --git a/zellij-server/src/plugins/plugin_map.rs b/zellij-server/src/plugins/plugin_map.rs
new file mode 100644
index 00000000..4cb5ebc9
--- /dev/null
+++ b/zellij-server/src/plugins/plugin_map.rs
@@ -0,0 +1,94 @@
+use crate::plugins::wasm_bridge::PluginId;
+use std::{
+    collections::{HashMap, HashSet},
+    path::PathBuf,
+    sync::{Arc, Mutex},
+};
+use wasmer::Instance;
+use wasmer_wasi::WasiEnv;
+
+use crate::{thread_bus::ThreadSenders, ClientId};
+
+use zellij_utils::{data::EventType, input::plugins::PluginConfig};
+
+// the idea here is to provide atomicity when adding/removing plugins from the map (eg. when a new
+// client connects) but to also allow updates/renders not to block each other
+// so when adding/removing from the map - everything is halted, that's life
+// but when cloning the internal RunningPlugin and Subscriptions atomics, we can call methods on
+// them without blocking other instances
+pub type PluginMap =
+    HashMap<(PluginId, ClientId), (Arc<Mutex<RunningPlugin>>, Arc<Mutex<Subscriptions>>)>;
+pub type Subscriptions = HashSet<EventType>;
+
+#[derive(Clone)]
+pub struct PluginEnv {
+    pub plugin_id: u32,
+    pub plugin: PluginConfig,
+    pub senders: ThreadSenders,
+    pub wasi_env: WasiEnv,
+    pub tab_index: usize,
+    pub client_id: ClientId,
+    #[allow(dead_code)]
+    pub plugin_own_data_dir: PathBuf,
+}
+
+impl PluginEnv {
+    // Get the name (path) of the containing plugin
+    pub fn name(&self) -> String {
+        format!(
+            "{} (ID {})",
+            self.plugin.path.display().to_string(),
+            self.plugin_id
+        )
+    }
+}
+
+#[derive(Eq, PartialEq, Hash)]
+pub enum AtomicEvent {
+    Resize,
+}
+
+pub struct RunningPlugin {
+    pub instance: Instance,
+    pub plugin_env: PluginEnv,
+    pub rows: usize,
+    pub columns: usize,
+    next_event_ids: HashMap<AtomicEvent, usize>,
+    last_applied_event_ids: HashMap<AtomicEvent, usize>,
+}
+
+impl RunningPlugin {
+    pub fn new(instance: Instance, plugin_env: PluginEnv, rows: usize, columns: usize) -> Self {
+        RunningPlugin {
+            instance,
+            plugin_env,
+            rows,
+            columns,
+            next_event_ids: HashMap::new(),
+            last_applied_event_ids: HashMap::new(),
+        }
+    }
+    pub fn next_event_id(&mut self, atomic_event: AtomicEvent) -> usize {
+        // TODO: probably not usize...
+        let current_event_id = *self.next_event_ids.get(&atomic_event).unwrap_or(&0);
+        if current_event_id < usize::MAX {
+            let next_event_id = current_event_id + 1;
+            self.next_event_ids.insert(atomic_event, next_event_id);
+            current_event_id
+        } else {
+            let current_event_id = 0;
+            let next_event_id = 1;
+            self.last_applied_event_ids.remove(&atomic_event);
+            self.next_event_ids.insert(atomic_event, next_event_id);
+            current_event_id
+        }
+    }
+    pub fn apply_event_id(&mut self, atomic_event: AtomicEvent, event_id: usize) -> bool {
+        if &event_id >= self.last_applied_event_ids.get(&atomic_event).unwrap_or(&0) {
+            self.last_applied_event_ids.insert(atomic_event, event_id);
+            true
+        } else {
+            false
+        }
+    }
+}
diff --git a/zellij-server/src/plugins/wasm_bridge.rs b/zellij-server/src/plugins/wasm_bridge.rs
index 119a1d24..45a516d3 100644
--- a/zellij-server/src/plugins/wasm_bridge.rs
+++ b/zellij-server/src/plugins/wasm_bridge.rs
@@ -1,79 +1,36 @@
 use super::PluginInstruction;
 use crate::plugins::plugin_loader::{PluginLoader, VersionMismatchError};
-use log::{debug, info, warn};
-use serde::{de::DeserializeOwned, Serialize};
+use crate::plugins::plugin_map::{AtomicEvent, PluginEnv, PluginMap};
+use crate::plugins::zellij_exports::{wasi_read_string, wasi_write_object};
+use log::info;
 use std::{
     collections::{HashMap, HashSet},
     fmt::Display,
     path::PathBuf,
-    process,
     str::FromStr,
     sync::{Arc, Mutex},
-    thread,
-    time::{Duration, Instant},
 };
-use wasmer::{
-    imports, ChainableNamedResolver, Function, ImportObject, Instance, Module, Store, Value,
-    WasmerEnv,
-};
-use wasmer_wasi::WasiEnv;
+use wasmer::{Instance, Module, Store, Value};
 use zellij_utils::async_std::task::{self, JoinHandle};
 
 use crate::{
-    background_jobs::BackgroundJob,
-    panes::PaneId,
-    pty::{ClientOrTabIndex, PtyInstruction},
-    screen::ScreenInstruction,
-    thread_bus::ThreadSenders,
-    ui::loading_indication::LoadingIndication,
-    ClientId,
+    background_jobs::BackgroundJob, screen::ScreenInstruction, thread_bus::ThreadSenders,
+    ui::loading_indication::LoadingIndication, ClientId,
 };
 
 use zellij_utils::{
     consts::VERSION,
-    data::{Event, EventType, PluginIds},
+    data::{Event, EventType},
     errors::prelude::*,
     errors::ZellijError,
     input::{
-        command::TerminalAction,
         layout::{RunPlugin, RunPluginLocation},
-        plugins::{PluginConfig, PluginType, PluginsConfig},
+        plugins::PluginsConfig,
     },
     pane_size::Size,
-    serde,
 };
 
-type PluginId = u32;
-
-#[derive(WasmerEnv, Clone)]
-pub struct PluginEnv {
-    pub plugin_id: u32,
-    pub plugin: PluginConfig,
-    pub senders: ThreadSenders,
-    pub wasi_env: WasiEnv,
-    pub subscriptions: Arc<Mutex<HashSet<EventType>>>,
-    pub tab_index: usize,
-    pub client_id: ClientId,
-    #[allow(dead_code)]
-    pub plugin_own_data_dir: PathBuf,
-}
-
-impl PluginEnv {
-    // Get the name (path) of the containing plugin
-    pub fn name(&self) -> String {
-        format!(
-            "{} (ID {})",
-            self.plugin.path.display().to_string(),
-            self.plugin_id
-        )
-    }
-}
-
-pub type PluginMap = HashMap<(u32, ClientId), (Instance, PluginEnv, (usize, usize))>; // u32 =>
-                                                                                      // plugin_id,
-                                                                                      // (usize, usize)
-                                                                                      // => (rows,
-                                                                                      // columns)
+pub type PluginId = u32;
 
 pub struct WasmBridge {
     connected_clients: Arc<Mutex<Vec<ClientId>>>,
@@ -121,10 +78,24 @@ impl WasmBridge {
         run: &RunPlugin,
         tab_index: usize,
         size: Size,
-        client_id: ClientId,
+        client_id: Option<ClientId>,
     ) -> Result<u32> {
         // returns the plugin id
-        let err_context = move || format!("failed to load plugin for client {client_id}");
+        let err_context = move || format!("failed to load plugin");
+
+        let client_id = client_id
+            .or_else(|| {
+                self.connected_clients
+                    .lock()
+                    .unwrap()
+                    .iter()
+                    .next()
+                    .copied()
+            })
+            .with_context(|| {
+                "Plugins must have a client id, none was provided and none are connected"
+            })?;
+
         let plugin_id = self.next_plugin_id;
 
         let plugin = self
@@ -277,48 +248,30 @@ impl WasmBridge {
         Ok(())
     }
     pub fn add_client(&mut self, client_id: ClientId) -> Result<()> {
-        let err_context = || format!("failed to add plugins for client {client_id}");
-
-        self.connected_clients.lock().unwrap().push(client_id);
-
-        let mut seen = HashSet::new();
-        let mut new_plugins = HashMap::new();
-        let mut plugin_map = self.plugin_map.lock().unwrap();
-        for (&(plugin_id, _), (instance, plugin_env, (rows, columns))) in &*plugin_map {
-            if seen.contains(&plugin_id) {
-                continue;
-            }
-            seen.insert(plugin_id);
-            let mut new_plugin_env = plugin_env.clone();
-
-            new_plugin_env.client_id = client_id;
-            new_plugins.insert(
-                plugin_id,
-                (instance.module().clone(), new_plugin_env, (*rows, *columns)),
-            );
+        let mut loading_indication = LoadingIndication::new("".into());
+        match PluginLoader::add_client(
+            client_id,
+            self.plugin_dir.clone(),
+            self.plugin_cache.clone(),
+            self.senders.clone(),
+            self.store.clone(),
+            self.plugin_map.clone(),
+            self.connected_clients.clone(),
+            &mut loading_indication,
+        ) {
+            Ok(_) => {
+                let _ = self
+                    .senders
+                    .send_to_screen(ScreenInstruction::RequestStateUpdateForPlugins);
+                Ok(())
+            },
+            Err(e) => Err(e),
         }
-        for (plugin_id, (module, mut new_plugin_env, (rows, columns))) in new_plugins.drain() {
-            let wasi = new_plugin_env
-                .wasi_env
-                .import_object(&module)
-                .with_context(err_context)?;
-            let zellij = zellij_exports(&self.store, &new_plugin_env);
-            let mut instance =
-                Instance::new(&module, &zellij.chain_back(wasi)).with_context(err_context)?;
-            load_plugin_instance(&mut instance).with_context(err_context)?;
-            plugin_map.insert(
-                (plugin_id, client_id),
-                (instance, new_plugin_env, (rows, columns)),
-            );
-        }
-        Ok(())
     }
     pub fn resize_plugin(&mut self, pid: u32, new_columns: usize, new_rows: usize) -> Result<()> {
-        let err_context = || format!("failed to resize plugin {pid}");
-        let mut plugin_bytes = vec![];
-        let mut plugin_map = self.plugin_map.lock().unwrap();
-        for ((plugin_id, client_id), (instance, plugin_env, (current_rows, current_columns))) in
-            plugin_map.iter_mut()
+        let err_context = move || format!("failed to resize plugin {pid}");
+        for ((plugin_id, client_id), (running_plugin, _subscriptions)) in
+            self.plugin_map.lock().unwrap().iter_mut()
         {
             if self
                 .cached_resizes_for_pending_plugins
@@ -327,26 +280,53 @@ impl WasmBridge {
                 continue;
             }
             if *plugin_id == pid {
-                *current_rows = new_rows;
-                *current_columns = new_columns;
-
-                // TODO: consolidate with above render function
-                let rendered_bytes = instance
-                    .exports
-                    .get_function("render")
-                    .map_err(anyError::new)
-                    .and_then(|render| {
-                        render
-                            .call(&[
-                                Value::I32(*current_rows as i32),
-                                Value::I32(*current_columns as i32),
-                            ])
-                            .map_err(anyError::new)
-                    })
-                    .and_then(|_| wasi_read_string(&plugin_env.wasi_env))
-                    .with_context(err_context)?;
-
-                plugin_bytes.push((*plugin_id, *client_id, rendered_bytes.as_bytes().to_vec()));
+                let event_id = running_plugin
+                    .lock()
+                    .unwrap()
+                    .next_event_id(AtomicEvent::Resize);
+                task::spawn({
+                    let senders = self.senders.clone();
+                    let running_plugin = running_plugin.clone();
+                    let plugin_id = *plugin_id;
+                    let client_id = *client_id;
+                    async move {
+                        let mut running_plugin = running_plugin.lock().unwrap();
+                        if running_plugin.apply_event_id(AtomicEvent::Resize, event_id) {
+                            running_plugin.rows = new_rows;
+                            running_plugin.columns = new_columns;
+                            let rendered_bytes = running_plugin
+                                .instance
+                                .exports
+                                .get_function("render")
+                                .map_err(anyError::new)
+                                .and_then(|render| {
+                                    render
+                                        .call(&[
+                                            Value::I32(running_plugin.rows as i32),
+                                            Value::I32(running_plugin.columns as i32),
+                                        ])
+                                        .map_err(anyError::new)
+                                })
+                                .and_then(|_| wasi_read_string(&running_plugin.plugin_env.wasi_env))
+                                .with_context(err_context);
+                            match rendered_bytes {
+                                Ok(rendered_bytes) => {
+                                    let plugin_bytes = vec![(
+                                        plugin_id,
+                                        client_id,
+                                        rendered_bytes.as_bytes().to_vec(),
+                                    )];
+                                    senders
+                                        .send_to_screen(ScreenInstruction::PluginBytes(
+                                            plugin_bytes,
+                                        ))
+                                        .unwrap();
+                                },
+                                Err(e) => log::error!("{}", e),
+                            }
+                        }
+                    }
+                });
             }
         }
         for (plugin_id, mut current_size) in self.cached_resizes_for_pending_plugins.iter_mut() {
@@ -355,9 +335,6 @@ impl WasmBridge {
                 current_size.1 = new_columns;
             }
         }
-        let _ = self
-            .senders
-            .send_to_screen(ScreenInstruction::PluginBytes(plugin_bytes));
         Ok(())
     }
     pub fn update_plugins(
@@ -366,21 +343,17 @@ impl WasmBridge {
     ) -> Result<()> {
         let err_context = || "failed to update plugin state".to_string();
 
-        let plugin_map = self.plugin_map.lock().unwrap();
-        let mut plugin_bytes = vec![];
         for (pid, cid, event) in updates.drain(..) {
-            for (&(plugin_id, client_id), (instance, plugin_env, (rows, columns))) in &*plugin_map {
+            for (&(plugin_id, client_id), (running_plugin, subscriptions)) in
+                &*self.plugin_map.lock().unwrap()
+            {
                 if self
                     .cached_events_for_pending_plugins
                     .contains_key(&plugin_id)
                 {
                     continue;
                 }
-                let subs = plugin_env
-                    .subscriptions
-                    .lock()
-                    .to_anyhow()
-                    .with_context(err_context)?;
+                let subs = subscriptions.lock().unwrap().clone();
                 // FIXME: This is very janky... Maybe I should write my own macro for Event -> EventType?
                 let event_type =
                     EventType::from_str(&event.to_string()).with_context(err_context)?;
@@ -390,16 +363,34 @@ impl WasmBridge {
                         || (cid.is_none() && pid == Some(plugin_id))
                         || (cid == Some(client_id) && pid == Some(plugin_id)))
                 {
-                    apply_event_to_plugin(
-                        plugin_id,
-                        client_id,
-                        &instance,
-                        &plugin_env,
-                        &event,
-                        *rows,
-                        *columns,
-                        &mut plugin_bytes,
-                    )?;
+                    task::spawn({
+                        let senders = self.senders.clone();
+                        let running_plugin = running_plugin.clone();
+                        let event = event.clone();
+                        async move {
+                            let running_plugin = running_plugin.lock().unwrap();
+                            let mut plugin_bytes = vec![];
+                            match apply_event_to_plugin(
+                                plugin_id,
+                                client_id,
+                                &running_plugin.instance,
+                                &running_plugin.plugin_env,
+                                &event,
+                                running_plugin.rows,
+                                running_plugin.columns,
+                                &mut plugin_bytes,
+                            ) {
+                                Ok(()) => {
+                                    let _ = senders.send_to_screen(ScreenInstruction::PluginBytes(
+                                        plugin_bytes,
+                                    ));
+                                },
+                                Err(e) => {
+                                    log::error!("{}", e);
+                                },
+                            }
+                        }
+                    });
                 }
             }
             for (plugin_id, cached_events) in self.cached_events_for_pending_plugins.iter_mut() {
@@ -408,9 +399,6 @@ impl WasmBridge {
                 }
             }
         }
-        let _ = self
-            .senders
-            .send_to_screen(ScreenInstruction::PluginBytes(plugin_bytes));
         Ok(())
     }
     pub fn apply_cached_events(&mut self, plugin_ids: Vec<u32>) -> Result<()> {
@@ -450,7 +438,6 @@ impl WasmBridge {
     fn apply_cached_events_and_resizes_for_plugin(&mut self, plugin_id: PluginId) -> Result<()> {
         let err_context = || format!("Failed to apply cached events to plugin");
         if let Some(events) = self.cached_events_for_pending_plugins.remove(&plugin_id) {
-            let mut plugin_map = self.plugin_map.lock().unwrap();
             let all_connected_clients: Vec<ClientId> = self
                 .connected_clients
                 .lock()
@@ -459,35 +446,48 @@ impl WasmBridge {
                 .copied()
                 .collect();
             for client_id in &all_connected_clients {
-                let mut plugin_bytes = vec![];
-                if let Some((instance, plugin_env, (rows, columns))) =
-                    plugin_map.get_mut(&(plugin_id, *client_id))
+                if let Some((running_plugin, subscriptions)) = self
+                    .plugin_map
+                    .lock()
+                    .unwrap()
+                    .get_mut(&(plugin_id, *client_id))
                 {
-                    let subs = plugin_env
-                        .subscriptions
-                        .lock()
-                        .to_anyhow()
-                        .with_context(err_context)?;
+                    let subs = subscriptions.lock().unwrap().clone();
                     for event in events.clone() {
                         let event_type =
                             EventType::from_str(&event.to_string()).with_context(err_context)?;
                         if !subs.contains(&event_type) {
                             continue;
                         }
-                        apply_event_to_plugin(
-                            plugin_id,
-                            *client_id,
-                            &instance,
-                            &plugin_env,
-                            &event,
-                            *rows,
-                            *columns,
-                            &mut plugin_bytes,
-                        )?;
+                        task::spawn({
+                            let senders = self.senders.clone();
+                            let running_plugin = running_plugin.clone();
+                            let client_id = *client_id;
+                            async move {
+                                let running_plugin = running_plugin.lock().unwrap();
+                                let mut plugin_bytes = vec![];
+                                match apply_event_to_plugin(
+                                    plugin_id,
+                                    client_id,
+                                    &running_plugin.instance,
+                                    &running_plugin.plugin_env,
+                                    &event,
+                                    running_plugin.rows,
+                                    running_plugin.columns,
+                                    &mut plugin_bytes,
+                                ) {
+                                    Ok(()) => {
+                                        let _ = senders.send_to_screen(
+                                            ScreenInstruction::PluginBytes(plugin_bytes),
+                                        );
+                                    },
+                                    Err(e) => {
+                                        log::error!("{}", e);
+                                    },
+                                }
+                            }
+                        });
                     }
-                    let _ = self
-                        .senders
-                        .send_to_screen(ScreenInstruction::PluginBytes(plugin_bytes));
                 }
             }
         }
@@ -512,11 +512,11 @@ impl WasmBridge {
             .lock()
             .unwrap()
             .iter()
-            .filter(
-                |((_plugin_id, _client_id), (_instance, plugin_env, _size))| {
-                    &plugin_env.plugin.location == plugin_location
-                },
-            )
+            .filter(|(_, (running_plugin, _subscriptions))| {
+                &running_plugin.lock().unwrap().plugin_env.plugin.location == plugin_location
+                // TODO:
+                // better
+            })
             .map(|((plugin_id, _client_id), _)| *plugin_id)
             .collect();
         if plugin_ids.is_empty() {
@@ -530,8 +530,11 @@ impl WasmBridge {
             .lock()
             .unwrap()
             .iter()
-            .find(|((p_id, _client_id), (_instance, _plugin_env, _size))| *p_id == plugin_id)
-            .map(|((_p_id, _client_id), (_instance, _plugin_env, size))| *size)
+            .find(|((p_id, _client_id), _)| *p_id == plugin_id)
+            .map(|(_, (running_plugin, _subscriptions))| {
+                let running_plugin = running_plugin.lock().unwrap();
+                (running_plugin.rows, running_plugin.columns)
+            })
     }
     fn start_plugin_loading_indication(
         &self,
@@ -563,6 +566,7 @@ fn handle_plugin_loading_failure(
     loading_indication: &mut LoadingIndication,
     error: impl Display,
 ) {
+    log::error!("{}", error);
     let _ = senders.send_to_background_jobs(BackgroundJob::StopPluginLoadingAnimation(plugin_id));
     loading_indication.indicate_loading_error(error.to_string());
     let _ = senders.send_to_screen(ScreenInstruction::UpdatePluginLoadingStage(
@@ -571,292 +575,6 @@ fn handle_plugin_loading_failure(
     ));
 }
 
-fn load_plugin_instance(instance: &mut Instance) -> Result<()> {
-    let err_context = || format!("failed to load plugin from instance {instance:#?}");
-
-    let load_function = instance
-        .exports
-        .get_function("_start")
-        .with_context(err_context)?;
-    // This eventually calls the `.load()` method
-    load_function.call(&[]).with_context(err_context)?;
-    Ok(())
-}
-
-pub(crate) fn zellij_exports(store: &Store, plugin_env: &PluginEnv) -> ImportObject {
-    macro_rules! zellij_export {
-        ($($host_function:ident),+ $(,)?) => {
-            imports! {
-                "zellij" => {
-                    $(stringify!($host_function) =>
-                        Function::new_native_with_env(store, plugin_env.clone(), $host_function),)+
-                }
-            }
-        }
-    }
-
-    zellij_export! {
-        host_subscribe,
-        host_unsubscribe,
-        host_set_selectable,
-        host_get_plugin_ids,
-        host_get_zellij_version,
-        host_open_file,
-        host_switch_tab_to,
-        host_set_timeout,
-        host_exec_cmd,
-        host_report_panic,
-    }
-}
-
-fn host_subscribe(plugin_env: &PluginEnv) {
-    wasi_read_object::<HashSet<EventType>>(&plugin_env.wasi_env)
-        .and_then(|new| {
-            plugin_env.subscriptions.lock().to_anyhow()?.extend(new);
-            Ok(())
-        })
-        .with_context(|| format!("failed to subscribe for plugin {}", plugin_env.name()))
-        .fatal();
-}
-
-fn host_unsubscribe(plugin_env: &PluginEnv) {
-    wasi_read_object::<HashSet<EventType>>(&plugin_env.wasi_env)
-        .and_then(|old| {
-            plugin_env
-                .subscriptions
-                .lock()
-                .to_anyhow()?
-                .retain(|k| !old.contains(k));
-            Ok(())
-        })
-        .with_context(|| format!("failed to unsubscribe for plugin {}", plugin_env.name()))
-        .fatal();
-}
-
-fn host_set_selectable(plugin_env: &PluginEnv, selectable: i32) {
-    match plugin_env.plugin.run {
-        PluginType::Pane(Some(tab_index)) => {
-            let selectable = selectable != 0;
-            plugin_env
-                .senders
-                .send_to_screen(ScreenInstruction::SetSelectable(
-                    PaneId::Plugin(plugin_env.plugin_id),
-                    selectable,
-                    tab_index,
-                ))
-                .with_context(|| {
-                    format!(
-                        "failed to set plugin {} selectable from plugin {}",
-                        selectable,
-                        plugin_env.name()
-                    )
-                })
-                .non_fatal();
-        },
-        _ => {
-            debug!(
-                "{} - Calling method 'host_set_selectable' does nothing for headless plugins",
-                plugin_env.plugin.location
-            )
-        },
-    }
-}
-
-fn host_get_plugin_ids(plugin_env: &PluginEnv) {
-    let ids = PluginIds {
-        plugin_id: plugin_env.plugin_id,
-        zellij_pid: process::id(),
-    };
-    wasi_write_object(&plugin_env.wasi_env, &ids)
-        .with_context(|| {
-            format!(
-                "failed to query plugin IDs from host for plugin {}",
-                plugin_env.name()
-            )
-        })
-        .non_fatal();
-}
-
-fn host_get_zellij_version(plugin_env: &PluginEnv) {
-    wasi_write_object(&plugin_env.wasi_env, VERSION)
-        .with_context(|| {
-            format!(
-                "failed to request zellij version from host for plugin {}",
-                plugin_env.name()
-            )
-        })
-        .non_fatal();
-}
-
-fn host_open_file(plugin_env: &PluginEnv) {
-    wasi_read_object::<PathBuf>(&plugin_env.wasi_env)
-        .and_then(|path| {
-            plugin_env
-                .senders
-                .send_to_pty(PtyInstruction::SpawnTerminal(
-                    Some(TerminalAction::OpenFile(path, None, None)),
-                    None,
-                    None,
-                    ClientOrTabIndex::TabIndex(plugin_env.tab_index),
-                ))
-        })
-        .with_context(|| {
-            format!(
-                "failed to open file on host from plugin {}",
-                plugin_env.name()
-            )
-        })
-        .non_fatal();
-}
-
-fn host_switch_tab_to(plugin_env: &PluginEnv, tab_idx: u32) {
-    plugin_env
-        .senders
-        .send_to_screen(ScreenInstruction::GoToTab(
-            tab_idx,
-            Some(plugin_env.client_id),
-        ))
-        .with_context(|| {
-            format!(
-                "failed to switch host to tab {tab_idx} from plugin {}",
-                plugin_env.name()
-            )
-        })
-        .non_fatal();
-}
-
-fn host_set_timeout(plugin_env: &PluginEnv, secs: f64) {
-    // There is a fancy, high-performance way to do this with zero additional threads:
-    // If the plugin thread keeps a BinaryHeap of timer structs, it can manage multiple and easily `.peek()` at the
-    // next time to trigger in O(1) time. Once the wake-up time is known, the `wasm` thread can use `recv_timeout()`
-    // to wait for an event with the timeout set to be the time of the next wake up. If events come in in the meantime,
-    // they are handled, but if the timeout triggers, we replace the event from `recv()` with an
-    // `Update(pid, TimerEvent)` and pop the timer from the Heap (or reschedule it). No additional threads for as many
-    // timers as we'd like.
-    //
-    // But that's a lot of code, and this is a few lines:
-    let send_plugin_instructions = plugin_env.senders.to_plugin.clone();
-    let update_target = Some(plugin_env.plugin_id);
-    let client_id = plugin_env.client_id;
-    let plugin_name = plugin_env.name();
-    thread::spawn(move || {
-        let start_time = Instant::now();
-        thread::sleep(Duration::from_secs_f64(secs));
-        // FIXME: The way that elapsed time is being calculated here is not exact; it doesn't take into account the
-        // time it takes an event to actually reach the plugin after it's sent to the `wasm` thread.
-        let elapsed_time = Instant::now().duration_since(start_time).as_secs_f64();
-
-        send_plugin_instructions
-            .ok_or(anyhow!("found no sender to send plugin instruction to"))
-            .and_then(|sender| {
-                sender
-                    .send(PluginInstruction::Update(vec![(
-                        update_target,
-                        Some(client_id),
-                        Event::Timer(elapsed_time),
-                    )]))
-                    .to_anyhow()
-            })
-            .with_context(|| {
-                format!(
-                    "failed to set host timeout of {secs} s for plugin {}",
-                    plugin_name
-                )
-            })
-            .non_fatal();
-    });
-}
-
-fn host_exec_cmd(plugin_env: &PluginEnv) {
-    let err_context = || {
-        format!(
-            "failed to execute command on host for plugin '{}'",
-            plugin_env.name()
-        )
-    };
-
-    let mut cmdline: Vec<String> = wasi_read_object(&plugin_env.wasi_env)
-        .with_context(err_context)
-        .fatal();
-    let command = cmdline.remove(0);
-
-    // Bail out if we're forbidden to run command
-    if !plugin_env.plugin._allow_exec_host_cmd {
-        warn!("This plugin isn't allow to run command in host side, skip running this command: '{cmd} {args}'.",
-        	cmd = command, args = cmdline.join(" "));
-        return;
-    }
-
-    // Here, we don't wait the command to finish
-    process::Command::new(command)
-        .args(cmdline)
-        .spawn()
-        .with_context(err_context)
-        .non_fatal();
-}
-
-// Custom panic handler for plugins.
-//
-// This is called when a panic occurs in a plugin. Since most panics will likely originate in the
-// code trying to deserialize an `Event` upon a plugin state update, we read some panic message,
-// formatted as string from the plugin.
-fn host_report_panic(plugin_env: &PluginEnv) {
-    let msg = wasi_read_string(&plugin_env.wasi_env)
-        .with_context(|| format!("failed to report panic for plugin '{}'", plugin_env.name()))
-        .fatal();
-    panic!("{}", msg);
-}
-
-// Helper Functions ---------------------------------------------------------------------------------------------------
-
-pub fn wasi_read_string(wasi_env: &WasiEnv) -> Result<String> {
-    let err_context = || format!("failed to read string from WASI env '{wasi_env:?}'");
-
-    let mut buf = String::new();
-    wasi_env
-        .state()
-        .fs
-        .stdout_mut()
-        .map_err(anyError::new)
-        .and_then(|stdout| {
-            stdout
-                .as_mut()
-                .ok_or(anyhow!("failed to get mutable reference to stdout"))
-        })
-        .and_then(|wasi_file| wasi_file.read_to_string(&mut buf).map_err(anyError::new))
-        .with_context(err_context)?;
-    // https://stackoverflow.com/questions/66450942/in-rust-is-there-a-way-to-make-literal-newlines-in-r-using-windows-c
-    Ok(buf.replace("\n", "\n\r"))
-}
-
-pub fn wasi_write_string(wasi_env: &WasiEnv, buf: &str) -> Result<()> {
-    wasi_env
-        .state()
-        .fs
-        .stdin_mut()
-        .map_err(anyError::new)
-        .and_then(|stdin| {
-            stdin
-                .as_mut()
-                .ok_or(anyhow!("failed to get mutable reference to stdin"))
-        })
-        .and_then(|stdin| writeln!(stdin, "{}\r", buf).map_err(anyError::new))
-        .with_context(|| format!("failed to write string to WASI env '{wasi_env:?}'"))
-}
-
-pub fn wasi_write_object(wasi_env: &WasiEnv, object: &(impl Serialize + ?Sized)) -> Result<()> {
-    serde_json::to_string(&object)
-        .map_err(anyError::new)
-        .and_then(|string| wasi_write_string(wasi_env, &string))
-        .with_context(|| format!("failed to serialize object for WASI env '{wasi_env:?}'"))
-}
-
-pub fn wasi_read_object<T: DeserializeOwned>(wasi_env: &WasiEnv) -> Result<T> {
-    wasi_read_string(wasi_env)
-        .and_then(|string| serde_json::from_str(&string).map_err(anyError::new))
-        .with_context(|| format!("failed to deserialize object from WASI env '{wasi_env:?}'"))
-}
-
 pub fn apply_event_to_plugin(
     plugin_id: u32,
     client_id: ClientId,
diff --git a/zellij-server/src/plugins/zellij_exports.rs b/zellij-server/src/plugins/zellij_exports.rs
new file mode 100644
index 00000000..cc0321ab
--- /dev/null
+++ b/zellij-server/src/plugins/zellij_exports.rs
@@ -0,0 +1,326 @@
+use super::PluginInstruction;
+use crate::plugins::plugin_map::{PluginEnv, Subscriptions};
+use log::{debug, warn};
+use serde::{de::DeserializeOwned, Serialize};
+use std::{
+    collections::HashSet,
+    path::PathBuf,
+    process,
+    sync::{Arc, Mutex},
+    thread,
+    time::{Duration, Instant},
+};
+use wasmer::{imports, Function, ImportObject, Store, WasmerEnv};
+use wasmer_wasi::WasiEnv;
+
+use crate::{
+    panes::PaneId,
+    pty::{ClientOrTabIndex, PtyInstruction},
+    screen::ScreenInstruction,
+};
+
+use zellij_utils::{
+    consts::VERSION,
+    data::{Event, EventType, PluginIds},
+    errors::prelude::*,
+    input::{command::TerminalAction, plugins::PluginType},
+    serde,
+};
+
+pub fn zellij_exports(
+    store: &Store,
+    plugin_env: &PluginEnv,
+    subscriptions: &Arc<Mutex<Subscriptions>>,
+) -> ImportObject {
+    macro_rules! zellij_export {
+        ($($host_function:ident),+ $(,)?) => {
+            imports! {
+                "zellij" => {
+                    $(stringify!($host_function) =>
+                        Function::new_native_with_env(store, ForeignFunctionEnv::new(plugin_env, subscriptions), $host_function),)+
+                }
+            }
+        }
+    }
+
+    zellij_export! {
+        host_subscribe,
+        host_unsubscribe,
+        host_set_selectable,
+        host_get_plugin_ids,
+        host_get_zellij_version,
+        host_open_file,
+        host_switch_tab_to,
+        host_set_timeout,
+        host_exec_cmd,
+        host_report_panic,
+    }
+}
+
+#[derive(WasmerEnv, Clone)]
+pub struct ForeignFunctionEnv {
+    pub plugin_env: PluginEnv,
+    pub subscriptions: Arc<Mutex<Subscriptions>>,
+}
+
+impl ForeignFunctionEnv {
+    pub fn new(plugin_env: &PluginEnv, subscriptions: &Arc<Mutex<Subscriptions>>) -> Self {
+        ForeignFunctionEnv {
+            plugin_env: plugin_env.clone(),
+            subscriptions: subscriptions.clone(),
+        }
+    }
+}
+
+fn host_subscribe(env: &ForeignFunctionEnv) {
+    wasi_read_object::<HashSet<EventType>>(&env.plugin_env.wasi_env)
+        .and_then(|new| {
+            env.subscriptions.lock().to_anyhow()?.extend(new);
+            Ok(())
+        })
+        .with_context(|| format!("failed to subscribe for plugin {}", env.plugin_env.name()))
+        .fatal();
+}
+
+fn host_unsubscribe(env: &ForeignFunctionEnv) {
+    wasi_read_object::<HashSet<EventType>>(&env.plugin_env.wasi_env)
+        .and_then(|old| {
+            env.subscriptions
+                .lock()
+                .to_anyhow()?
+                .retain(|k| !old.contains(k));
+            Ok(())
+        })
+        .with_context(|| format!("failed to unsubscribe for plugin {}", env.plugin_env.name()))
+        .fatal();
+}
+
+fn host_set_selectable(env: &ForeignFunctionEnv, selectable: i32) {
+    match env.plugin_env.plugin.run {
+        PluginType::Pane(Some(tab_index)) => {
+            let selectable = selectable != 0;
+            env.plugin_env
+                .senders
+                .send_to_screen(ScreenInstruction::SetSelectable(
+                    PaneId::Plugin(env.plugin_env.plugin_id),
+                    selectable,
+                    tab_index,
+                ))
+                .with_context(|| {
+                    format!(
+                        "failed to set plugin {} selectable from plugin {}",
+                        selectable,
+                        env.plugin_env.name()
+                    )
+                })
+                .non_fatal();
+        },
+        _ => {
+            debug!(
+                "{} - Calling method 'host_set_selectable' does nothing for headless plugins",
+                env.plugin_env.plugin.location
+            )
+        },
+    }
+}
+
+fn host_get_plugin_ids(env: &ForeignFunctionEnv) {
+    let ids = PluginIds {
+        plugin_id: env.plugin_env.plugin_id,
+        zellij_pid: process::id(),
+    };
+    wasi_write_object(&env.plugin_env.wasi_env, &ids)
+        .with_context(|| {
+            format!(
+                "failed to query plugin IDs from host for plugin {}",
+                env.plugin_env.name()
+            )
+        })
+        .non_fatal();
+}
+
+fn host_get_zellij_version(env: &ForeignFunctionEnv) {
+    wasi_write_object(&env.plugin_env.wasi_env, VERSION)
+        .with_context(|| {
+            format!(
+                "failed to request zellij version from host for plugin {}",
+                env.plugin_env.name()
+            )
+        })
+        .non_fatal();
+}
+
+fn host_open_file(env: &ForeignFunctionEnv) {
+    wasi_read_object::<PathBuf>(&env.plugin_env.wasi_env)
+        .and_then(|path| {
+            env.plugin_env
+                .senders
+                .send_to_pty(PtyInstruction::SpawnTerminal(
+                    Some(TerminalAction::OpenFile(path, None, None)),
+                    None,
+                    None,
+                    ClientOrTabIndex::TabIndex(env.plugin_env.tab_index),
+                ))
+        })
+        .with_context(|| {
+            format!(
+                "failed to open file on host from plugin {}",
+                env.plugin_env.name()
+            )
+        })
+        .non_fatal();
+}
+
+fn host_switch_tab_to(env: &ForeignFunctionEnv, tab_idx: u32) {
+    env.plugin_env
+        .senders
+        .send_to_screen(ScreenInstruction::GoToTab(
+            tab_idx,
+            Some(env.plugin_env.client_id),
+        ))
+        .with_context(|| {
+            format!(
+                "failed to switch host to tab {tab_idx} from plugin {}",
+                env.plugin_env.name()
+            )
+        })
+        .non_fatal();
+}
+
+fn host_set_timeout(env: &ForeignFunctionEnv, secs: f64) {
+    // There is a fancy, high-performance way to do this with zero additional threads:
+    // If the plugin thread keeps a BinaryHeap of timer structs, it can manage multiple and easily `.peek()` at the
+    // next time to trigger in O(1) time. Once the wake-up time is known, the `wasm` thread can use `recv_timeout()`
+    // to wait for an event with the timeout set to be the time of the next wake up. If events come in in the meantime,
+    // they are handled, but if the timeout triggers, we replace the event from `recv()` with an
+    // `Update(pid, TimerEvent)` and pop the timer from the Heap (or reschedule it). No additional threads for as many
+    // timers as we'd like.
+    //
+    // But that's a lot of code, and this is a few lines:
+    let send_plugin_instructions = env.plugin_env.senders.to_plugin.clone();
+    let update_target = Some(env.plugin_env.plugin_id);
+    let client_id = env.plugin_env.client_id;
+    let plugin_name = env.plugin_env.name();
+    thread::spawn(move || {
+        let start_time = Instant::now();
+        thread::sleep(Duration::from_secs_f64(secs));
+        // FIXME: The way that elapsed time is being calculated here is not exact; it doesn't take into account the
+        // time it takes an event to actually reach the plugin after it's sent to the `wasm` thread.
+        let elapsed_time = Instant::now().duration_since(start_time).as_secs_f64();
+
+        send_plugin_instructions
+            .ok_or(anyhow!("found no sender to send plugin instruction to"))
+            .and_then(|sender| {
+                sender
+                    .send(PluginInstruction::Update(vec![(
+                        update_target,
+                        Some(client_id),
+                        Event::Timer(elapsed_time),
+                    )]))
+                    .to_anyhow()
+            })
+            .with_context(|| {
+                format!(
+                    "failed to set host timeout of {secs} s for plugin {}",
+                    plugin_name
+                )
+            })
+            .non_fatal();
+    });
+}
+
+fn host_exec_cmd(env: &ForeignFunctionEnv) {
+    let err_context = || {
+        format!(
+            "failed to execute command on host for plugin '{}'",
+            env.plugin_env.name()
+        )
+    };
+
+    let mut cmdline: Vec<String> = wasi_read_object(&env.plugin_env.wasi_env)
+        .with_context(err_context)
+        .fatal();
+    let command = cmdline.remove(0);
+
+    // Bail out if we're forbidden to run command
+    if !env.plugin_env.plugin._allow_exec_host_cmd {
+        warn!("This plugin isn't allow to run command in host side, skip running this command: '{cmd} {args}'.",
+        	cmd = command, args = cmdline.join(" "));
+        return;
+    }
+
+    // Here, we don't wait the command to finish
+    process::Command::new(command)
+        .args(cmdline)
+        .spawn()
+        .with_context(err_context)
+        .non_fatal();
+}
+
+// Custom panic handler for plugins.
+//
+// This is called when a panic occurs in a plugin. Since most panics will likely originate in the
+// code trying to deserialize an `Event` upon a plugin state update, we read some panic message,
+// formatted as string from the plugin.
+fn host_report_panic(env: &ForeignFunctionEnv) {
+    let msg = wasi_read_string(&env.plugin_env.wasi_env)
+        .with_context(|| {
+            format!(
+                "failed to report panic for plugin '{}'",
+                env.plugin_env.name()
+            )
+        })
+        .fatal();
+    panic!("{}", msg);
+}
+
+// Helper Functions ---------------------------------------------------------------------------------------------------
+
+pub fn wasi_read_string(wasi_env: &WasiEnv) -> Result<String> {
+    let err_context = || format!("failed to read string from WASI env '{wasi_env:?}'");
+
+    let mut buf = vec![];
+    wasi_env
+        .state()
+        .fs
+        .stdout_mut()
+        .map_err(anyError::new)
+        .and_then(|stdout| {
+            stdout
+                .as_mut()
+                .ok_or(anyhow!("failed to get mutable reference to stdout"))
+        })
+        .and_then(|wasi_file| wasi_file.read_to_end(&mut buf).map_err(anyError::new))
+        .with_context(err_context)?;
+    let buf = String::from_utf8_lossy(&buf);
+    // https://stackoverflow.com/questions/66450942/in-rust-is-there-a-way-to-make-literal-newlines-in-r-using-windows-c
+    Ok(buf.replace("\n", "\n\r"))
+}
+
+pub fn wasi_write_string(wasi_env: &WasiEnv, buf: &str) -> Result<()> {
+    wasi_env
+        .state()
+        .fs
+        .stdin_mut()
+        .map_err(anyError::new)
+        .and_then(|stdin| {
+            stdin
+                .as_mut()
+                .ok_or(anyhow!("failed to get mutable reference to stdin"))
+        })
+        .and_then(|stdin| writeln!(stdin, "{}\r", buf).map_err(anyError::new))
+        .with_context(|| format!("failed to write string to WASI env '{wasi_env:?}'"))
+}
+
+pub fn wasi_write_object(wasi_env: &WasiEnv, object: &(impl Serialize + ?Sized)) -> Result<()> {
+    serde_json::to_string(&object)
+        .map_err(anyError::new)
+        .and_then(|string| wasi_write_string(wasi_env, &string))
+        .with_context(|| format!("failed to serialize object for WASI env '{wasi_env:?}'"))
+}
+
+pub fn wasi_read_object<T: DeserializeOwned>(wasi_env: &WasiEnv) -> Result<T> {
+    wasi_read_string(wasi_env)
+        .and_then(|string| serde_json::from_str(&string).map_err(anyError::new))
+        .with_context(|| format!("failed to deserialize object from WASI env '{wasi_env:?}'"))
+}