lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0

//! WASM plugin runtime using wasmi.
//!
//! This module provides the core WASM execution environment for plugins.
//! It handles module loading, instantiation, and function calls.

#[cfg(feature = "wasm-plugins")]
use wasmi::{
    Caller, Config, Engine, Extern, Func, Instance, Linker, Memory, MemoryType, Module, Store,
    TypedFunc,
};

use alloc::string::{String, ToString};
use alloc::vec;
use alloc::vec::Vec;

use super::error::PluginError;
use super::host::{HostFunctions, HostState, LogLevel};
use super::types::{PluginContext, PluginLimits, PluginManifest};

// ═══════════════════════════════════════════════════════════════════════════════
// WASM PLUGIN
// ═══════════════════════════════════════════════════════════════════════════════

/// A loaded WASM plugin instance.
#[cfg(feature = "wasm-plugins")]
pub struct WasmPlugin {
    /// Plugin name
    name: String,
    /// WASM engine
    engine: Engine,
    /// WASM store with host state
    store: Store<HostState>,
    /// WASM instance
    instance: Instance,
    /// Plugin manifest
    manifest: PluginManifest,
    /// Invocation count
    invocation_count: u64,
    /// Total execution time in microseconds
    total_execution_us: u64,
}

#[cfg(feature = "wasm-plugins")]
impl WasmPlugin {
    /// Load a WASM plugin from bytes.
    pub fn load(name: &str, wasm_bytes: &[u8]) -> Result<Self, PluginError> {
        Self::load_with_limits(name, wasm_bytes, PluginLimits::default())
    }

    /// Load a WASM plugin with custom resource limits.
    pub fn load_with_limits(
        name: &str,
        wasm_bytes: &[u8],
        limits: PluginLimits,
    ) -> Result<Self, PluginError> {
        // Create engine with configuration
        let mut config = Config::default();
        config.consume_fuel(true); // Enable fuel metering for time limits

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

        // Compile module
        let module = Module::new(&engine, wasm_bytes)
            .map_err(|e| PluginError::CompilationFailed(alloc::format!("{:?}", e)))?;

        // Create store with host state
        let host_state = HostState {
            limits,
            ..Default::default()
        };
        let mut store = Store::new(&engine, host_state);

        // Set fuel limit based on time limit (rough approximation)
        // ~1000 fuel per millisecond
        let fuel = store.data().limits.max_execution_ms * 1000;
        store.set_fuel(fuel).ok();

        // Create linker and add host functions
        let mut linker: Linker<HostState> = Linker::new(&engine);
        Self::register_host_functions(&mut linker)?;

        // Instantiate module
        let instance = linker
            .instantiate(&mut store, &module)
            .map_err(|e| PluginError::InstantiationFailed(alloc::format!("{:?}", e)))?
            .start(&mut store)
            .map_err(|e| PluginError::InstantiationFailed(alloc::format!("{:?}", e)))?;

        // Call plugin_init
        let init_result = Self::call_init(&instance, &mut store)?;
        if init_result != 0 {
            return Err(PluginError::InitFailed(init_result));
        }

        // Get manifest
        let manifest = Self::read_manifest(&instance, &mut store)?;

        // Update host state with manifest
        store.data_mut().manifest = manifest.clone();

        Ok(Self {
            name: name.into(),
            engine,
            store,
            instance,
            manifest,
            invocation_count: 0,
            total_execution_us: 0,
        })
    }

    /// Register host functions with the linker.
    fn register_host_functions(linker: &mut Linker<HostState>) -> Result<(), PluginError> {
        // host_log(level: i32, msg_ptr: i32, msg_len: i32)
        linker
            .func_wrap(
                "env",
                "host_log",
                |mut caller: Caller<HostState>, level: i32, msg_ptr: i32, msg_len: i32| {
                    if let Some(memory) = get_memory(&caller) {
                        if let Ok(msg) = read_string(&memory, &caller, msg_ptr, msg_len) {
                            HostFunctions::host_log(caller.data_mut(), level, &msg);
                        }
                    }
                },
            )
            .map_err(|e| PluginError::Internal(alloc::format!("linker error: {:?}", e)))?;

        // host_get_config(key_ptr: i32, key_len: i32, out_ptr: i32, out_cap: i32) -> i32
        linker
            .func_wrap(
                "env",
                "host_get_config",
                |mut caller: Caller<HostState>,
                 key_ptr: i32,
                 key_len: i32,
                 out_ptr: i32,
                 out_cap: i32|
                 -> i32 {
                    let memory = match get_memory(&caller) {
                        Some(m) => m,
                        None => return -1,
                    };

                    let key = match read_string(&memory, &caller, key_ptr, key_len) {
                        Ok(k) => k,
                        Err(_) => return -1,
                    };

                    // Get config value
                    let value = match caller.data().config.get(&key) {
                        Some(v) => v.clone(),
                        None => return -1,
                    };

                    // Write to output buffer
                    let bytes = value.as_bytes();
                    let copy_len = bytes.len().min(out_cap as usize);
                    if write_bytes(&memory, &mut caller, out_ptr, &bytes[..copy_len]).is_err() {
                        return -1;
                    }

                    copy_len as i32
                },
            )
            .map_err(|e| PluginError::Internal(alloc::format!("linker error: {:?}", e)))?;

        // host_read_file(path_ptr: i32, path_len: i32, out_ptr: i32, out_cap: i32) -> i32
        linker
            .func_wrap(
                "env",
                "host_read_file",
                |mut caller: Caller<HostState>,
                 path_ptr: i32,
                 path_len: i32,
                 out_ptr: i32,
                 out_cap: i32|
                 -> i32 {
                    let memory = match get_memory(&caller) {
                        Some(m) => m,
                        None => return -1,
                    };

                    let path = match read_string(&memory, &caller, path_ptr, path_len) {
                        Ok(p) => p,
                        Err(_) => return -1,
                    };

                    // Read file via host state
                    let mut buf = vec![0u8; out_cap as usize];
                    match HostFunctions::host_read_file(caller.data(), &path, &mut buf) {
                        Ok(len) => {
                            if write_bytes(&memory, &mut caller, out_ptr, &buf[..len]).is_err() {
                                return -1;
                            }
                            len as i32
                        }
                        Err(_) => -1,
                    }
                },
            )
            .map_err(|e| PluginError::Internal(alloc::format!("linker error: {:?}", e)))?;

        // host_file_exists(path_ptr: i32, path_len: i32) -> i32
        linker
            .func_wrap(
                "env",
                "host_file_exists",
                |caller: Caller<HostState>, path_ptr: i32, path_len: i32| -> i32 {
                    let memory = match get_memory(&caller) {
                        Some(m) => m,
                        None => return -1,
                    };

                    let path = match read_string(&memory, &caller, path_ptr, path_len) {
                        Ok(p) => p,
                        Err(_) => return -1,
                    };

                    HostFunctions::host_file_exists(caller.data(), &path)
                },
            )
            .map_err(|e| PluginError::Internal(alloc::format!("linker error: {:?}", e)))?;

        // host_file_size(path_ptr: i32, path_len: i32) -> i64
        linker
            .func_wrap(
                "env",
                "host_file_size",
                |caller: Caller<HostState>, path_ptr: i32, path_len: i32| -> i64 {
                    let memory = match get_memory(&caller) {
                        Some(m) => m,
                        None => return -1,
                    };

                    let path = match read_string(&memory, &caller, path_ptr, path_len) {
                        Ok(p) => p,
                        Err(_) => return -1,
                    };

                    HostFunctions::host_file_size(caller.data(), &path)
                },
            )
            .map_err(|e| PluginError::Internal(alloc::format!("linker error: {:?}", e)))?;

        Ok(())
    }

    /// Call plugin_init and return result code.
    fn call_init(instance: &Instance, store: &mut Store<HostState>) -> Result<i32, PluginError> {
        let init: TypedFunc<(), i32> = instance
            .get_typed_func(store, "plugin_init")
            .map_err(|_| PluginError::MissingExport("plugin_init".into()))?;

        init.call(store, ())
            .map_err(|e| PluginError::ExecutionFailed(alloc::format!("{:?}", e)))
    }

    /// Read plugin manifest from WASM module.
    fn read_manifest(
        instance: &Instance,
        store: &mut Store<HostState>,
    ) -> Result<PluginManifest, PluginError> {
        // Get plugin_manifest function
        let manifest_fn: TypedFunc<(i32, i32), i32> = instance
            .get_typed_func(store, "plugin_manifest")
            .map_err(|_| PluginError::MissingExport("plugin_manifest".into()))?;

        // Allocate buffer in WASM memory for manifest
        let memory = instance
            .get_memory(store, "memory")
            .ok_or_else(|| PluginError::MemoryError("no memory export".into()))?;

        // Use a fixed buffer in linear memory
        // First, check if alloc function exists
        let alloc_result = instance.get_typed_func::<i32, i32>(store, "alloc");

        let (manifest_ptr, manifest_cap) = if let Ok(alloc) = alloc_result {
            // Plugin provides alloc function
            let ptr = alloc.call(store, 4096).map_err(|e| {
                PluginError::ExecutionFailed(alloc::format!("alloc failed: {:?}", e))
            })?;
            (ptr, 4096)
        } else {
            // Use fixed address (plugins without alloc must reserve this space)
            (0x10000, 4096)
        };

        // Call plugin_manifest
        let len = manifest_fn
            .call(store, (manifest_ptr, manifest_cap))
            .map_err(|e| PluginError::ExecutionFailed(alloc::format!("{:?}", e)))?;

        if len < 0 {
            return Err(PluginError::ExecutionFailed(
                "plugin_manifest returned error".into(),
            ));
        }

        // Read manifest JSON from WASM memory
        let mut buf = vec![0u8; len as usize];
        memory
            .read(store, manifest_ptr as usize, &mut buf)
            .map_err(|e| PluginError::MemoryError(alloc::format!("{:?}", e)))?;

        let json = String::from_utf8(buf).map_err(|_| {
            PluginError::InvalidManifest(super::types::ManifestError::InvalidFormat)
        })?;

        PluginManifest::from_json(&json).map_err(PluginError::InvalidManifest)
    }

    /// Get plugin name.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Get plugin manifest.
    pub fn manifest(&self) -> &PluginManifest {
        &self.manifest
    }

    /// Get host state (mutable).
    pub fn host_state_mut(&mut self) -> &mut HostState {
        self.store.data_mut()
    }

    /// Get host state (immutable).
    pub fn host_state(&self) -> &HostState {
        self.store.data()
    }

    /// Process data through the plugin.
    ///
    /// This calls the plugin's `plugin_process` function.
    pub fn process(&mut self, ctx: &PluginContext, data: &[u8]) -> Result<Vec<u8>, PluginError> {
        // Check fuel/time limit
        if self.store.get_fuel().unwrap_or(0) == 0 {
            return Err(PluginError::ResourceLimitExceeded("execution time".into()));
        }

        // Get memory and alloc function
        let memory = self
            .instance
            .get_memory(&self.store, "memory")
            .ok_or_else(|| PluginError::MemoryError("no memory export".into()))?;

        // Check memory limit
        let mem_pages = memory.size(&self.store);
        let mem_bytes = mem_pages as usize * 65536; // 64KB per page
        if mem_bytes > self.store.data().limits.max_memory {
            return Err(PluginError::ResourceLimitExceeded("memory".into()));
        }

        // Serialize context to JSON
        let ctx_json = ctx.to_json();
        let ctx_bytes = ctx_json.as_bytes();

        // Calculate output capacity (2x input for potential expansion)
        let out_cap = (data.len() * 2).max(4096);
        if out_cap > self.store.data().limits.max_output_size {
            return Err(PluginError::ResourceLimitExceeded("output size".into()));
        }

        // Allocate memory in WASM
        let total_needed = ctx_bytes.len() + data.len() + out_cap;

        let (ctx_ptr, data_ptr, out_ptr) = if let Ok(alloc) = self
            .instance
            .get_typed_func::<i32, i32>(&self.store, "alloc")
        {
            let ctx_ptr = alloc
                .call(&mut self.store, ctx_bytes.len() as i32)
                .map_err(|e| PluginError::ExecutionFailed(alloc::format!("{:?}", e)))?;
            let data_ptr = alloc
                .call(&mut self.store, data.len() as i32)
                .map_err(|e| PluginError::ExecutionFailed(alloc::format!("{:?}", e)))?;
            let out_ptr = alloc
                .call(&mut self.store, out_cap as i32)
                .map_err(|e| PluginError::ExecutionFailed(alloc::format!("{:?}", e)))?;
            (ctx_ptr, data_ptr, out_ptr)
        } else {
            // Fixed layout for plugins without alloc
            let ctx_ptr = 0x10000i32;
            let data_ptr = ctx_ptr + ctx_bytes.len() as i32 + 16;
            let out_ptr = data_ptr + data.len() as i32 + 16;
            (ctx_ptr, data_ptr, out_ptr)
        };

        // Write data to WASM memory
        memory
            .write(&mut self.store, ctx_ptr as usize, ctx_bytes)
            .map_err(|e| PluginError::MemoryError(alloc::format!("{:?}", e)))?;
        memory
            .write(&mut self.store, data_ptr as usize, data)
            .map_err(|e| PluginError::MemoryError(alloc::format!("{:?}", e)))?;

        // Call plugin_process
        let process: TypedFunc<(i32, i32, i32, i32, i32, i32), i32> = self
            .instance
            .get_typed_func(&self.store, "plugin_process")
            .map_err(|_| PluginError::MissingExport("plugin_process".into()))?;

        let result_len = process
            .call(
                &mut self.store,
                (
                    ctx_ptr,
                    ctx_bytes.len() as i32,
                    data_ptr,
                    data.len() as i32,
                    out_ptr,
                    out_cap as i32,
                ),
            )
            .map_err(|e| {
                // Check if it's a fuel exhaustion
                if alloc::format!("{:?}", e).contains("fuel") {
                    PluginError::ResourceLimitExceeded("execution time".into())
                } else {
                    PluginError::ExecutionFailed(alloc::format!("{:?}", e))
                }
            })?;

        if result_len < 0 {
            return Err(PluginError::PluginReturnedError(result_len));
        }

        // Read output
        let mut output = vec![0u8; result_len as usize];
        memory
            .read(&self.store, out_ptr as usize, &mut output)
            .map_err(|e| PluginError::MemoryError(alloc::format!("{:?}", e)))?;

        // Update stats
        self.invocation_count += 1;

        Ok(output)
    }

    /// Cleanup plugin resources.
    pub fn destroy(&mut self) -> Result<(), PluginError> {
        // Call plugin_destroy if it exists
        if let Ok(destroy) = self
            .instance
            .get_typed_func::<(), ()>(&self.store, "plugin_destroy")
        {
            destroy
                .call(&mut self.store, ())
                .map_err(|e| PluginError::ExecutionFailed(alloc::format!("{:?}", e)))?;
        }
        Ok(())
    }

    /// Get invocation count.
    pub fn invocation_count(&self) -> u64 {
        self.invocation_count
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// HELPER FUNCTIONS
// ═══════════════════════════════════════════════════════════════════════════════

#[cfg(feature = "wasm-plugins")]
fn get_memory<T>(caller: &Caller<T>) -> Option<Memory> {
    caller.get_export("memory")?.into_memory()
}

#[cfg(feature = "wasm-plugins")]
fn read_string<T>(
    memory: &Memory,
    caller: &Caller<T>,
    ptr: i32,
    len: i32,
) -> Result<String, PluginError> {
    let mut buf = vec![0u8; len as usize];
    memory
        .read(caller, ptr as usize, &mut buf)
        .map_err(|e| PluginError::MemoryError(alloc::format!("{:?}", e)))?;
    String::from_utf8(buf).map_err(|_| PluginError::MemoryError("invalid utf8".into()))
}

#[cfg(feature = "wasm-plugins")]
fn write_bytes<T>(
    memory: &Memory,
    caller: &mut Caller<T>,
    ptr: i32,
    data: &[u8],
) -> Result<(), PluginError> {
    memory
        .write(caller, ptr as usize, data)
        .map_err(|e| PluginError::MemoryError(alloc::format!("{:?}", e)))
}

// ═══════════════════════════════════════════════════════════════════════════════
// NO-OP IMPLEMENTATION (when wasm-plugins feature is disabled)
// ═══════════════════════════════════════════════════════════════════════════════

/// Stub plugin when WASM feature is disabled.
#[cfg(not(feature = "wasm-plugins"))]
pub struct WasmPlugin {
    name: String,
    manifest: PluginManifest,
}

#[cfg(not(feature = "wasm-plugins"))]
impl WasmPlugin {
    /// Always returns error when WASM is disabled.
    pub fn load(_name: &str, _wasm_bytes: &[u8]) -> Result<Self, PluginError> {
        Err(PluginError::Internal(
            "WASM plugins feature not enabled".into(),
        ))
    }

    /// Always returns error when WASM is disabled.
    pub fn load_with_limits(
        _name: &str,
        _wasm_bytes: &[u8],
        _limits: PluginLimits,
    ) -> Result<Self, PluginError> {
        Err(PluginError::Internal(
            "WASM plugins feature not enabled".into(),
        ))
    }

    /// Get plugin name.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Get plugin manifest.
    pub fn manifest(&self) -> &PluginManifest {
        &self.manifest
    }

    /// Get host state (mutable) - stub version returns panic.
    pub fn host_state_mut(&mut self) -> &mut HostState {
        unreachable!("WASM plugins feature not enabled")
    }

    /// Always returns error.
    pub fn process(&mut self, _ctx: &PluginContext, _data: &[u8]) -> Result<Vec<u8>, PluginError> {
        Err(PluginError::Internal(
            "WASM plugins feature not enabled".into(),
        ))
    }

    /// No-op.
    pub fn destroy(&mut self) -> Result<(), PluginError> {
        Ok(())
    }

    /// Always zero.
    pub fn invocation_count(&self) -> u64 {
        0
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// TESTS
// ═══════════════════════════════════════════════════════════════════════════════

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    #[cfg(not(feature = "wasm-plugins"))]
    fn test_wasm_disabled() {
        let result = WasmPlugin::load("test", &[]);
        assert!(result.is_err());
    }
}