mimium-lang 4.0.0-alpha

// Wasmtime-based WASM runtime
//
// This module provides execution of WASM modules generated by the WASM backend.
// It uses Wasmtime as the WASM runtime and bridges RuntimePrimitives to host functions.

pub mod engine;

use crate::runtime::primitives::Word;
use crate::runtime::vm::heap::{self, HeapStorage};
use std::collections::HashMap;
use wasmtime::{
    AsContextMut, Caller, Config, Engine, FuncType, Linker, Module, OptLevel, Store, Val, ValType,
    WasmBacktraceDetails,
};

/// Upper bound of `$arityN` builtin specializations exported by the WASM host runtime.
///
/// Must match compiler-side import declarations in `compiler/wasmgen.rs`.
const MAX_SPECIALIZED_BUILTIN_ARITY: usize = 16;

/// WASM-side analogue of [`SystemPluginAudioWorker`](crate::plugin::SystemPluginAudioWorker).
///
/// Plugins that need per-sample processing on the WASM backend implement
/// this trait on their audio-thread handle.  [`WasmDspRuntime`] holds a
/// `Vec<Box<dyn WasmSystemPluginAudioWorker>>` and calls [`on_sample`]
/// for each worker before invoking `dsp()`.
pub trait WasmSystemPluginAudioWorker: Send {
    /// Called once per sample, before `dsp()`.
    ///
    /// The worker receives the current time and a mutable reference to the
    /// [`WasmEngine`] so that it can execute WASM-side closures (e.g. via
    /// `_mimium_exec_closure_void`).
    fn on_sample(
        &mut self,
        time: crate::runtime::Time,
        engine: &mut engine::WasmEngine,
    ) -> crate::runtime::vm::ReturnCode;
}

/// A single plugin function callable from WASM trampolines.
pub type WasmPluginFn = std::sync::Arc<dyn Fn(&[f64]) -> Option<f64> + Send + Sync>;

/// A map of plugin function names to their implementations.
///
/// Plugin authors create this map from their audio handle (e.g. via
/// `GuiAudioHandle::into_wasm_plugin_fn_map`). Each entry maps a function
/// name to a closure that handles the call.
pub type WasmPluginFnMap = HashMap<String, WasmPluginFn>;

/// WASM runtime state
pub struct WasmRuntime {
    /// Wasmtime engine
    engine: Engine,
    /// Linker for connecting host functions
    linker: Linker<RuntimeState>,
}

/// State storage for a single execution context (global or per-closure).
/// Mirrors the native VM's `StateStorage` struct.
#[derive(Debug, Clone, Default)]
struct StateStorage {
    /// Current read/write cursor in the data array
    pos: usize,
    /// Flat array of state words (delay buffers, mem values, etc.)
    data: Vec<u64>,
}

impl StateStorage {
    fn with_size(size: usize) -> Self {
        Self {
            pos: 0,
            data: vec![0u64; size],
        }
    }
}

/// Per-instance runtime state passed to host functions
pub struct RuntimeState {
    /// Linear memory pointer (set after instantiation)
    memory: Option<wasmtime::Memory>,
    /// Heap storage for dynamic allocations
    heap: HeapStorage,
    /// Array storage (array ID -> array data)
    arrays: HashMap<Word, Vec<Word>>,
    /// Global state storage (used at the top-level/global context)
    global_state: StateStorage,
    /// Per-closure state storages, keyed by closure address in linear memory.
    /// Each closure gets its own StateStorage, allocated lazily on first call.
    closure_states: HashMap<i64, StateStorage>,
    /// Stack of closure addresses for tracking the active state context.
    /// When empty, global_state is active. When non-empty, the top entry's
    /// closure state is active. Mirrors the native VM's `states_stack`.
    state_stack: Vec<i64>,
    /// Current time (for runtime_get_now)
    pub(crate) current_time: u64,
    /// Sample rate (for runtime_get_samplerate)
    pub(crate) sample_rate: f64,
}

impl RuntimeState {
    /// Get the currently active state storage.
    /// Returns the global state if no closure is on the state stack,
    /// otherwise returns the state of the closure at the top of the stack.
    /// This mirrors the native VM's `get_current_state()` method.
    fn get_current_state(&mut self) -> &mut StateStorage {
        if let Some(&closure_addr) = self.state_stack.last() {
            self.closure_states
                .get_mut(&closure_addr)
                .expect("closure_state_push must be called before accessing closure state")
        } else {
            &mut self.global_state
        }
    }
}

impl Default for RuntimeState {
    fn default() -> Self {
        Self {
            memory: None,
            heap: HeapStorage::default(),
            arrays: HashMap::new(),
            global_state: StateStorage::default(),
            closure_states: HashMap::new(),
            state_stack: Vec::new(),
            current_time: 0,
            sample_rate: 44100.0,
        }
    }
}

impl WasmRuntime {
    /// Create a new WASM runtime with JIT compilation.
    ///
    /// `ext_fns` provides the complete set of external function type info
    /// so that plugin host trampolines can be registered for all plugins
    /// (system and dynamic).
    /// `plugin_fns` is an optional map of plugin function name to handler
    /// closures (see `WasmPluginFnMap`).
    #[cfg(not(target_arch = "wasm32"))]
    pub fn new(
        ext_fns: &[crate::plugin::ExtFunTypeInfo],
        plugin_fns: Option<WasmPluginFnMap>,
    ) -> Result<Self, String> {
        // Configure Wasmtime with JIT compiler and optimizations
        let mut config = Config::new();

        // Enable Cranelift JIT compiler with optimization level Speed
        config.cranelift_opt_level(OptLevel::Speed);

        // Enable parallel compilation for faster module loading
        config.parallel_compilation(true);

        // Enable WASM features that may improve performance
        config.wasm_simd(true); // SIMD operations
        config.wasm_bulk_memory(true); // Bulk memory operations
        config.wasm_backtrace_details(WasmBacktraceDetails::Enable);
        config.debug_info(true);

        let engine =
            Engine::new(&config).map_err(|e| format!("Failed to create WASM engine: {e}"))?;
        let mut linker = Linker::new(&engine);

        // Register all runtime primitive host functions
        Self::register_runtime_primitives(&mut linker)?;

        // Register plugin functions as host trampolines
        Self::register_plugin_functions(&mut linker, ext_fns, plugin_fns)?;

        Ok(Self { engine, linker })
    }

    /// Create a new WASM runtime for wasm32 target (without plugins)
    #[cfg(target_arch = "wasm32")]
    pub fn new(ext_fns: &[crate::plugin::ExtFunTypeInfo]) -> Result<Self, String> {
        // Configure Wasmtime with JIT compiler and optimizations
        let mut config = Config::new();

        // Enable Cranelift JIT compiler with optimization level Speed
        config.cranelift_opt_level(OptLevel::Speed);

        // Enable parallel compilation for faster module loading
        config.parallel_compilation(true);

        // Enable WASM features that may improve performance
        config.wasm_simd(true); // SIMD operations
        config.wasm_bulk_memory(true); // Bulk memory operations
        config.wasm_backtrace_details(WasmBacktraceDetails::Enable);
        config.debug_info(true);

        let engine =
            Engine::new(&config).map_err(|e| format!("Failed to create WASM engine: {e}"))?;
        let mut linker = Linker::new(&engine);

        // Register all runtime primitive host functions
        Self::register_runtime_primitives(&mut linker)?;

        // Register plugin functions as host trampolines (no plugin_fns on wasm32)
        Self::register_plugin_functions(&mut linker, ext_fns, None)?;

        Ok(Self { engine, linker })
    }

    /// Load and instantiate a WASM module
    pub fn load_module(&mut self, wasm_bytes: &[u8]) -> Result<WasmModule, String> {
        let module = Module::from_binary(&self.engine, wasm_bytes)
            .map_err(|e| format!("Failed to load WASM module: {e:#}"))?;

        self.instantiate_module(module)
    }

    /// Load and instantiate a precompiled/serialized WASM module artifact.
    pub fn load_precompiled_module(
        &mut self,
        serialized_module: &[u8],
    ) -> Result<WasmModule, String> {
        let module = unsafe { Module::deserialize(&self.engine, serialized_module) }
            .map_err(|e| format!("Failed to deserialize precompiled WASM module: {e:#}"))?;

        self.instantiate_module(module)
    }

    fn instantiate_module(&mut self, module: Module) -> Result<WasmModule, String> {
        let mut runtime_state = RuntimeState::default();

        let mut store = Store::new(&self.engine, runtime_state);

        let instance = self
            .linker
            .instantiate(&mut store, &module)
            .map_err(|e| format!("Failed to instantiate WASM module: {e:#}"))?;

        // Get memory export
        let memory = instance
            .get_memory(&mut store, "memory")
            .ok_or("WASM module does not export memory")?;

        store.data_mut().memory = Some(memory);

        Ok(WasmModule {
            module,
            store,
            instance,
            function_cache: std::collections::HashMap::new(),
        })
    }

    /// Register all runtime primitive host functions
    fn register_runtime_primitives(linker: &mut Linker<RuntimeState>) -> Result<(), String> {
        macro_rules! register {
            ($($name:expr => $func:expr),* $(,)?) => {
                $(
                    linker
                        .func_wrap("runtime", $name, $func)
                        .map_err(|e| format!("Failed to register {}: {}", $name, e))?;
                )*
            };
        }

        register! {
            // Heap operations
            "heap_alloc" => heap_alloc_host,
            "heap_retain" => heap_retain_host,
            "heap_release" => heap_release_host,
            "heap_load" => heap_load_host,
            "heap_store" => heap_store_host,
            // Box operations
            "box_alloc" => box_alloc_host,
            "box_load" => box_load_host,
            "box_clone" => box_clone_host,
            "box_release" => box_release_host,
            "box_store" => box_store_host,
            // UserSum operations
            "usersum_clone" => usersum_clone_host,
            "usersum_release" => usersum_release_host,
            // Closure operations
            "closure_make" => closure_make_host,
            "closure_close" => closure_close_host,
            "closure_call" => closure_call_host,
            // Closure state stack operations (per-closure state isolation)
            "closure_state_push" => closure_state_push_host,
            "closure_state_pop" => closure_state_pop_host,
            // State operations
            "state_push" => state_push_host,
            "state_pop" => state_pop_host,
            "state_get" => state_get_host,
            "state_set" => state_set_host,
            "state_delay" => state_delay_host,
            "state_mem" => state_mem_host,
            // Array operations
            "array_alloc" => array_alloc_host,
            "array_get_elem" => array_get_elem_host,
            "array_set_elem" => array_set_elem_host,
            // Runtime globals
            "runtime_get_now" => runtime_get_now_host,
            "runtime_get_samplerate" => runtime_get_samplerate_host,
        }

        // Register math functions (from "math" module)
        macro_rules! register_math_f1 {
            ($($name:expr => $func:expr),* $(,)?) => {
                $(
                    linker
                        .func_wrap("math", $name, $func)
                        .map_err(|e| format!("Failed to register math::{}: {}", $name, e))?;
                )*
            };
        }

        macro_rules! register_math_f2 {
            ($($name:expr => $func:expr),* $(,)?) => {
                $(
                    linker
                        .func_wrap("math", $name, $func)
                        .map_err(|e| format!("Failed to register math::{}: {}", $name, e))?;
                )*
            };
        }

        register_math_f1! {
            "sin" => |_caller: Caller<'_, RuntimeState>, x: f64| -> f64 { x.sin() },
            "cos" => |_caller: Caller<'_, RuntimeState>, x: f64| -> f64 { x.cos() },
            "tan" => |_caller: Caller<'_, RuntimeState>, x: f64| -> f64 { x.tan() },
            "sinh" => |_caller: Caller<'_, RuntimeState>, x: f64| -> f64 { x.sinh() },
            "cosh" => |_caller: Caller<'_, RuntimeState>, x: f64| -> f64 { x.cosh() },
            "tanh" => |_caller: Caller<'_, RuntimeState>, x: f64| -> f64 { x.tanh() },
            "asin" => |_caller: Caller<'_, RuntimeState>, x: f64| -> f64 { x.asin() },
            "acos" => |_caller: Caller<'_, RuntimeState>, x: f64| -> f64 { x.acos() },
            "atan" => |_caller: Caller<'_, RuntimeState>, x: f64| -> f64 { x.atan() },
            "round" => |_caller: Caller<'_, RuntimeState>, x: f64| -> f64 { x.round() },
            "floor" => |_caller: Caller<'_, RuntimeState>, x: f64| -> f64 { x.floor() },
            "ceil" => |_caller: Caller<'_, RuntimeState>, x: f64| -> f64 { x.ceil() },
            "log" => |_caller: Caller<'_, RuntimeState>, x: f64| -> f64 { x.ln() },
        }

        register_math_f2! {
            "atan2" => |_caller: Caller<'_, RuntimeState>, y: f64, x: f64| -> f64 { y.atan2(x) },
            "pow" => |_caller: Caller<'_, RuntimeState>, base: f64, exp: f64| -> f64 { base.powf(exp) },
            "min" => |_caller: Caller<'_, RuntimeState>, a: f64, b: f64| -> f64 { a.min(b) },
            "max" => |_caller: Caller<'_, RuntimeState>, a: f64, b: f64| -> f64 { a.max(b) },
        }

        // Register builtin functions (from "builtin" module)
        macro_rules! register_builtin {
            ($($name:expr => $func:expr),* $(,)?) => {
                $(
                    linker
                        .func_wrap("builtin", $name, $func)
                        .map_err(|e| format!("Failed to register builtin::{}: {}", $name, e))?;
                )*
            };
        }

        register_builtin! {
            "probeln" => builtin_probeln_host,
            "probe" => builtin_probe_host,
            "length_array" => builtin_length_array_host,
            "split_head" => builtin_split_head_host,
            "split_tail" => builtin_split_tail_host,
        }

        linker
            .func_new(
                "builtin",
                "prepend",
                FuncType::new(
                    linker.engine(),
                    vec![ValType::I64, ValType::I64],
                    vec![ValType::I64],
                ),
                move |caller, params, results| {
                    builtin_prepend_arity_host(caller, params, results, 1);
                    Ok(())
                },
            )
            .map_err(|e| format!("Failed to register builtin::prepend: {e}"))?;

        for arity in 1..=MAX_SPECIALIZED_BUILTIN_ARITY {
            let split_head_name = format!("split_head$arity{arity}");
            linker
                .func_wrap(
                    "builtin",
                    split_head_name.as_str(),
                    move |caller: Caller<'_, RuntimeState>, array: i64, dst_ptr: i32| {
                        builtin_split_head_arity_host(caller, array, dst_ptr, arity)
                    },
                )
                .map_err(|e| format!("Failed to register builtin::split_head$arity{arity}: {e}"))?;

            let split_tail_name = format!("split_tail$arity{arity}");
            linker
                .func_wrap(
                    "builtin",
                    split_tail_name.as_str(),
                    move |caller: Caller<'_, RuntimeState>, array: i64, dst_ptr: i32| {
                        builtin_split_tail_arity_host(caller, array, dst_ptr, arity)
                    },
                )
                .map_err(|e| format!("Failed to register builtin::split_tail$arity{arity}: {e}"))?;

            let prepend_name = format!("prepend$arity{arity}");
            linker
                .func_new(
                    "builtin",
                    prepend_name.as_str(),
                    FuncType::new(
                        linker.engine(),
                        vec![ValType::I64, ValType::I64],
                        vec![ValType::I64],
                    ),
                    move |caller, params, results| {
                        builtin_prepend_arity_host(caller, params, results, arity);
                        Ok(())
                    },
                )
                .map_err(|e| format!("Failed to register builtin::prepend$arity{arity}: {e}"))?;
        }

        Ok(())
    }

    /// Register plugin functions as host functions.
    ///
    /// Dynamically generates trampolines for all plugin functions based on their
    /// type info. The WASM function signature is derived from each plugin's
    /// `ExtFunTypeInfo`, matching what `wasmgen::setup_plugin_imports` produces.
    /// Register plugin host functions so that the WASM module can call them.
    ///
    /// For each plugin function we create a lightweight trampoline that
    /// inspects the WASM-level arguments and returns a sensible default.
    /// The rule is simple: **return the first argument unchanged** for any
    /// function whose return type matches its first parameter type (the
    /// typical "intercept / passthrough" pattern used by Probe), and return
    /// a zero otherwise.
    ///
    /// **Note:** This function only registers trampolines based on the provided
    /// `ext_fns`. Dynamic plugin loading is handled by the caller (`ExecContext`)
    /// to avoid duplicate registration.
    #[cfg(not(target_arch = "wasm32"))]
    fn register_plugin_functions(
        linker: &mut Linker<RuntimeState>,
        ext_fns: &[crate::plugin::ExtFunTypeInfo],
        plugin_fns: Option<WasmPluginFnMap>,
    ) -> Result<(), String> {
        use crate::types::Type;

        // Register trampolines for every runtime-stage external function.
        for type_info in ext_fns {
            let name = type_info.name;
            let fn_ty = type_info.ty.to_type();

            let Type::Function { arg, ret } = fn_ty else {
                log::warn!("Plugin '{}' has non-function type, skipping", name.as_str());
                continue;
            };

            // Build wasmtime FuncType matching wasmgen's setup_plugin_imports.
            // Flatten tuple arguments into separate WASM parameters
            let mut param_valtypes: Vec<ValType> = match arg.to_type() {
                Type::Tuple(elems) => {
                    // Flatten tuple arguments
                    elems
                        .iter()
                        .map(|t| mimium_type_to_wasmtime_valtype(&t.to_type()))
                        .collect()
                }
                arg_ty => {
                    // Single argument
                    vec![mimium_type_to_wasmtime_valtype(&arg_ty)]
                }
            };

            let ret_type = ret.to_type();
            let is_void = matches!(ret_type, Type::Primitive(crate::types::PType::Unit));
            let flat_ret_count = flat_return_word_count(&ret_type);
            let uses_dest_ptr =
                !is_void && flat_ret_count > 1 && plugin_uses_dest_ptr(name.as_str());

            let return_valtypes: Vec<ValType> = if uses_dest_ptr {
                param_valtypes.push(ValType::I32); // dst_ptr
                vec![] // void return — result written to memory
            } else if is_void {
                vec![]
            } else {
                vec![mimium_type_to_wasmtime_valtype(&ret_type)]
            };
            let plugin_name = name.as_str().to_string();

            log::debug!(
                "Registering plugin host function: {} ({} params -> {:?}{})",
                plugin_name,
                param_valtypes.len(),
                return_valtypes,
                if uses_dest_ptr { " [dest-ptr]" } else { "" },
            );

            let func_type = FuncType::new(
                linker.engine(),
                param_valtypes.clone(),
                return_valtypes.clone(),
            );

            let return_vts = return_valtypes.clone();

            // Look up a specific handler closure for this function, if any.
            let handler: Option<WasmPluginFn> = plugin_fns
                .as_ref()
                .and_then(|map| map.get(name.as_str()).cloned());

            if uses_dest_ptr {
                // Dest-pointer trampoline: the host writes multi-word results
                // directly into WASM linear memory at a caller-provided address.
                let ret_words = flat_ret_count;
                linker
                    .func_new(
                        "plugin",
                        name.as_str(),
                        func_type,
                        move |mut caller, params, _results| {
                            log::trace!(
                                "Plugin dest-ptr trampoline called: {plugin_name}({params:?})"
                            );

                            // Last parameter is the I32 destination pointer.
                            let dst_ptr = match params.last() {
                                Some(wasmtime::Val::I32(p)) => *p as usize,
                                _ => {
                                    log::error!(
                                        "{plugin_name}: missing dest pointer in last param"
                                    );
                                    return Ok(());
                                }
                            };

                            // Extract f64 value arguments (everything before dst_ptr).
                            let value_params = &params[..params.len() - 1];
                            let args: Vec<f64> = value_params
                                .iter()
                                .filter_map(|v| match v {
                                    wasmtime::Val::F64(f) => Some(f64::from_bits(*f)),
                                    wasmtime::Val::I64(i) => Some(*i as f64),
                                    wasmtime::Val::I32(i) => Some(*i as f64),
                                    _ => None,
                                })
                                .collect();

                            // Call the handler for side effects (e.g. push to
                            // ring buffers).
                            if let Some(ref func) = handler {
                                let _ = func(&args);
                            }

                            // Write passthrough values to the destination in
                            // WASM linear memory.  The probe intercept pattern
                            // returns its value arguments unchanged.
                            let memory = caller
                                .get_export("memory")
                                .and_then(|e| e.into_memory())
                                .expect("WASM module must export 'memory'");
                            let mem_data = memory.data_mut(&mut caller);
                            for i in 0..ret_words.min(args.len()) {
                                let offset = dst_ptr + i * 8;
                                let end = offset + 8;
                                if end <= mem_data.len() {
                                    mem_data[offset..end].copy_from_slice(&args[i].to_le_bytes());
                                }
                            }

                            Ok(())
                        },
                    )
                    .map_err(|e| format!("Failed to register plugin '{}': {e}", name.as_str()))?;
            } else {
                // Standard trampoline (single-word or void return).
                // Determine if this function is a "passthrough" shape (first param
                // type == return type, name contains "intercept").
                let is_passthrough = !return_valtypes.is_empty()
                    && param_valtypes.len() >= 2
                    && param_valtypes
                        .first()
                        .is_some_and(|first| valtype_eq(first, &return_valtypes[0]))
                    && plugin_name.contains("intercept");

                linker
                    .func_new(
                        "plugin",
                        name.as_str(),
                        func_type,
                        move |_caller, params, results| {
                            log::trace!("Plugin trampoline called: {plugin_name}({params:?})");

                            // If the function returns void (Unit), results is empty.
                            // Still call the handler for its side effects, but don't
                            // write any return value.
                            if results.is_empty() {
                                if let Some(ref func) = handler {
                                    let args = decode_trampoline_args(params);
                                    let _ = func(&args);
                                }
                                return Ok(());
                            }

                            // Try the per-function handler registered by the plugin
                            if let Some(ref func) = handler {
                                let args = decode_trampoline_args(params);

                                if let Some(result) = func(&args) {
                                    // Match the Val variant to the declared return type.
                                    // For multi-result signatures, preserve passthrough behavior
                                    // by copying remaining result slots from corresponding params.
                                    results[0] = match &return_vts[0] {
                                        ValType::I64 => wasmtime::Val::I64(result as i64),
                                        ValType::I32 => wasmtime::Val::I32(result as i32),
                                        _ => wasmtime::Val::F64(result.to_bits()),
                                    };
                                    if results.len() > 1 {
                                        for i in 1..results.len() {
                                            if let Some(param) = params.get(i) {
                                                results[i] = *param;
                                            } else {
                                                results[i] =
                                                    default_val_for_valtype(return_vts[i].clone());
                                            }
                                        }
                                    }
                                    return Ok(());
                                }
                            }

                            // Fallback: generic trampoline behavior
                            if is_passthrough && !params.is_empty() {
                                for (i, result) in results.iter_mut().enumerate() {
                                    *result = params.get(i).copied().unwrap_or_else(|| {
                                        default_val_for_valtype(return_vts[i].clone())
                                    });
                                }
                            } else {
                                for (i, result) in results.iter_mut().enumerate() {
                                    *result = default_val_for_valtype(return_vts[i].clone());
                                }
                            }
                            Ok(())
                        },
                    )
                    .map_err(|e| format!("Failed to register plugin '{}': {e}", name.as_str()))?;
            }
        }

        Ok(())
    }
}

impl Default for WasmRuntime {
    fn default() -> Self {
        Self::new(&[], None).expect("Failed to create WASM runtime")
    }
}

/// WASM module instance
pub struct WasmModule {
    #[allow(dead_code)]
    module: Module,
    store: Store<RuntimeState>,
    instance: wasmtime::Instance,
    /// Cache of frequently called functions (e.g., dsp)
    function_cache: std::collections::HashMap<String, wasmtime::Func>,
}

impl WasmModule {
    /// Serialize compiled module artifacts for fast later deserialization.
    pub fn serialize_compiled_module(&self) -> Result<Vec<u8>, String> {
        self.module
            .serialize()
            .map_err(|e| format!("Failed to serialize compiled WASM module: {e:#}"))
    }

    /// Get or cache a function by name
    pub fn get_or_cache_function(&mut self, name: &str) -> Result<wasmtime::Func, String> {
        if let Some(func) = self.function_cache.get(name) {
            return Ok(*func);
        }

        let func = self
            .instance
            .get_func(&mut self.store, name)
            .ok_or_else(|| format!("Function '{name}' not found in WASM module"))?;

        self.function_cache.insert(name.to_string(), func);
        Ok(func)
    }

    /// Call a function directly using a cached Func handle (faster)
    pub fn call_func_direct(
        &mut self,
        func: &wasmtime::Func,
        args: &[Word],
    ) -> Result<Vec<Word>, String> {
        // Convert args to wasmtime values
        let wasm_args: Vec<wasmtime::Val> =
            args.iter().map(|&w| wasmtime::Val::I64(w as i64)).collect();

        // Prepare results buffer
        let mut results = vec![wasmtime::Val::I64(0); func.ty(&self.store).results().len()];

        // Call function
        func.call(&mut self.store, &wasm_args, &mut results)
            .map_err(|e| {
                // Include detailed error chain
                format!("Failed to call function: {e:#}")
            })?;

        // Convert results back to Words
        Ok(results
            .into_iter()
            .map(|v| match v {
                wasmtime::Val::I64(i) => i as u64,
                wasmtime::Val::F64(f) => f,
                wasmtime::Val::I32(i) => i as u64,
                wasmtime::Val::F32(f) => f as u64,
                _ => 0,
            })
            .collect())
    }

    /// Call a function exported by the WASM module
    pub fn call_function(&mut self, name: &str, args: &[Word]) -> Result<Vec<Word>, String> {
        let func = self.get_or_cache_function(name)?;
        self.call_func_direct(&func, args)
    }

    /// Read an f64 value from linear memory at the given byte offset
    pub fn read_memory_f64(&mut self, offset: usize) -> Result<f64, String> {
        let memory = self
            .instance
            .get_memory(&mut self.store, "memory")
            .ok_or("No memory export")?;
        let data = memory.data(&self.store);
        if offset + 8 > data.len() {
            return Err(format!(
                "Memory read out of bounds: offset={offset}, memory size={}",
                data.len()
            ));
        }
        let bytes: [u8; 8] = data[offset..offset + 8]
            .try_into()
            .map_err(|e| format!("Failed to read memory: {e}"))?;
        Ok(f64::from_le_bytes(bytes))
    }

    /// Get mutable access to the runtime state
    /// This allows external code to update current_time, sample_rate, etc.
    pub fn get_runtime_state_mut(&mut self) -> Option<&mut RuntimeState> {
        Some(self.store.data_mut())
    }

    /// Update the current time in the runtime state.
    pub fn set_current_time(&mut self, time: u64) {
        self.store.data_mut().current_time = time;
    }

    /// Read current value of exported allocator pointer global.
    pub fn get_alloc_ptr(&mut self) -> Result<i32, String> {
        let global = self
            .instance
            .get_global(&mut self.store, "__alloc_ptr")
            .ok_or("No __alloc_ptr global export")?;
        match global.get(&mut self.store) {
            wasmtime::Val::I32(v) => Ok(v),
            _ => Err("__alloc_ptr global type mismatch".to_string()),
        }
    }

    /// Restore allocator pointer global to a previous value.
    pub fn set_alloc_ptr(&mut self, value: i32) -> Result<(), String> {
        let global = self
            .instance
            .get_global(&mut self.store, "__alloc_ptr")
            .ok_or("No __alloc_ptr global export")?;
        global
            .set(&mut self.store, wasmtime::Val::I32(value))
            .map_err(|e| format!("Failed to set __alloc_ptr: {e:#}"))
    }
}

/// Check whether the named external function uses the destination-pointer
/// calling convention, where the host writes multi-word results directly
/// into WASM linear memory instead of returning them normally.
///
/// This must stay in sync with `WasmGenerator::ext_function_uses_dest_ptr`
/// in `wasmgen.rs`.
#[cfg(not(target_arch = "wasm32"))]
fn plugin_uses_dest_ptr(name: &str) -> bool {
    name.starts_with("split_head")
        || name.starts_with("split_tail")
        || name.starts_with("__probe_intercept$arity")
        || name.starts_with("__probe_value_intercept$arity")
}

/// Count how many flat f64 words a mimium type expands to when stored in
/// WASM linear memory.  Single-word types return 1; tuples/records return
/// the sum of their element word counts.
#[cfg(not(target_arch = "wasm32"))]
fn flat_return_word_count(ty: &crate::types::Type) -> usize {
    use crate::types::{PType, Type};
    match ty {
        Type::Primitive(PType::Unit) => 0,
        Type::Tuple(elems) => elems
            .iter()
            .map(|e| flat_return_word_count(&e.to_type()))
            .sum(),
        Type::Record(fields) => fields
            .iter()
            .map(|f| flat_return_word_count(&f.ty.to_type()))
            .sum(),
        _ => 1,
    }
}

/// Convert a mimium `Type` to a `wasmtime::ValType`.
///
/// This mapping must stay in sync with `WasmGenerator::type_to_valtype` in
/// `wasmgen.rs` so that the host trampolines match the WASM import signatures.
#[cfg(not(target_arch = "wasm32"))]
fn mimium_type_to_wasmtime_valtype(ty: &crate::types::Type) -> ValType {
    use crate::types::{PType, Type};
    match ty {
        Type::Primitive(PType::Numeric) => ValType::F64,
        Type::Primitive(PType::Int) => ValType::I64,
        Type::Primitive(PType::String) => ValType::I64,
        Type::Primitive(PType::Unit) => ValType::I64,
        Type::Function { .. } => ValType::I64,
        Type::Record(fields) if fields.len() == 1 => {
            mimium_type_to_wasmtime_valtype(&fields[0].ty.to_type())
        }
        Type::Tuple(elems) if elems.len() == 1 => {
            mimium_type_to_wasmtime_valtype(&elems[0].to_type())
        }
        Type::Tuple(_) | Type::Record(_) => ValType::I64,
        Type::Array(_) => ValType::I64,
        Type::Union(_) | Type::UserSum { .. } => ValType::I64,
        Type::Ref(_) => ValType::I64,
        Type::Boxed(_) => ValType::I64,
        Type::Code(_) => ValType::I64,
        Type::Intermediate(cell) => {
            let tv = cell.read().unwrap();
            tv.parent.as_ref().map_or(ValType::I64, |parent| {
                mimium_type_to_wasmtime_valtype(&parent.to_type())
            })
        }
        _ => ValType::I64,
    }
}

/// Produce a zero/default `Val` for a given `ValType`.
/// Compare two [`ValType`] values for equality.
///
/// `wasmtime::ValType` does not implement `PartialEq`, so we match on the
/// discriminant instead.
#[cfg(not(target_arch = "wasm32"))]
fn valtype_eq(a: &ValType, b: &ValType) -> bool {
    matches!(
        (a, b),
        (ValType::I32, ValType::I32)
            | (ValType::I64, ValType::I64)
            | (ValType::F32, ValType::F32)
            | (ValType::F64, ValType::F64)
    )
}

#[cfg(not(target_arch = "wasm32"))]
fn decode_trampoline_args(params: &[Val]) -> Vec<f64> {
    params
        .iter()
        .filter_map(|v| match v {
            Val::F64(f) => Some(f64::from_bits(*f)),
            Val::I64(i) => Some(*i as f64),
            Val::I32(i) => Some(*i as f64),
            _ => None,
        })
        .collect()
}

#[cfg(not(target_arch = "wasm32"))]
fn default_val_for_valtype(vt: ValType) -> Val {
    match vt {
        ValType::F64 => Val::F64(0.0f64.to_bits()),
        ValType::F32 => Val::F32(0.0f32.to_bits()),
        ValType::I64 => Val::I64(0),
        ValType::I32 => Val::I32(0),
        _ => Val::I64(0),
    }
}

// Host function implementations

fn heap_alloc_host(mut caller: Caller<'_, RuntimeState>, size_words: i32) -> i64 {
    log::trace!("heap_alloc_host: size_words={size_words}");
    let state = caller.data_mut();
    let heap_obj = heap::HeapObject::new(size_words as usize);
    let heap_idx = state.heap.insert(heap_obj);

    // Convert HeapIdx to Word for WASM
    // We use the slotmap key's internal representation
    unsafe { std::mem::transmute::<heap::HeapIdx, u64>(heap_idx) as i64 }
}

fn heap_retain_host(mut caller: Caller<'_, RuntimeState>, obj: i64) {
    log::trace!("heap_retain_host: obj={obj}");
    let state = caller.data_mut();
    let heap_idx: heap::HeapIdx = unsafe { std::mem::transmute::<u64, heap::HeapIdx>(obj as u64) };
    heap::heap_retain(&mut state.heap, heap_idx);
}

fn heap_release_host(mut caller: Caller<'_, RuntimeState>, obj: i64) {
    log::trace!("heap_release_host: obj={obj}");
    let state = caller.data_mut();
    let heap_idx: heap::HeapIdx = unsafe { std::mem::transmute::<u64, heap::HeapIdx>(obj as u64) };
    heap::heap_release(&mut state.heap, heap_idx);
}

fn heap_load_host(mut caller: Caller<'_, RuntimeState>, dst_ptr: i32, obj: i64, size_words: i32) {
    log::trace!("heap_load_host: dst_ptr={dst_ptr}, obj={obj}, size_words={size_words}");

    let heap_idx: heap::HeapIdx = unsafe { std::mem::transmute::<u64, heap::HeapIdx>(obj as u64) };
    let size = size_words as usize;

    // Get heap data
    let data = {
        let state = caller.data();
        state
            .heap
            .get(heap_idx)
            .map(|heap_obj| heap_obj.data[..size].to_vec())
            .expect("heap_load: invalid heap index")
    };

    // Write to linear memory
    let memory = caller.data().memory.expect("Memory not initialized");
    let offset = dst_ptr as usize;
    let bytes = unsafe {
        std::slice::from_raw_parts(
            data.as_ptr() as *const u8,
            size * std::mem::size_of::<Word>(),
        )
    };
    memory
        .write(&mut caller, offset, bytes)
        .expect("Failed to write to WASM memory");
}

fn heap_store_host(mut caller: Caller<'_, RuntimeState>, obj: i64, src_ptr: i32, size_words: i32) {
    log::trace!("heap_store_host: obj={obj}, src_ptr={src_ptr}, size_words={size_words}");

    let heap_idx: heap::HeapIdx = unsafe { std::mem::transmute::<u64, heap::HeapIdx>(obj as u64) };
    let size = size_words as usize;

    // Read from linear memory
    let mut buffer = vec![0u64; size];
    let memory = caller.data().memory.expect("Memory not initialized");
    let offset = src_ptr as usize;
    let bytes = unsafe {
        std::slice::from_raw_parts_mut(
            buffer.as_mut_ptr() as *mut u8,
            size * std::mem::size_of::<Word>(),
        )
    };
    memory
        .read(&caller, offset, bytes)
        .expect("Failed to read from WASM memory");

    // Write to heap
    let state = caller.data_mut();
    let heap_obj = state
        .heap
        .get_mut(heap_idx)
        .expect("heap_store: invalid heap index");
    heap_obj.data[..size].copy_from_slice(&buffer[..size]);
}

fn box_alloc_host(mut caller: Caller<'_, RuntimeState>, src_ptr: i32, size_words: i32) -> i64 {
    log::trace!("box_alloc_host: src_ptr={src_ptr}, size_words={size_words}");

    let size = size_words as usize;

    // Read from linear memory
    let mut buffer = vec![0u64; size];
    let memory = caller.data().memory.expect("Memory not initialized");
    let offset = src_ptr as usize;
    let bytes = unsafe {
        std::slice::from_raw_parts_mut(
            buffer.as_mut_ptr() as *mut u8,
            size * std::mem::size_of::<Word>(),
        )
    };
    memory
        .read(&caller, offset, bytes)
        .expect("Failed to read from WASM memory");

    // Create heap object
    let state = caller.data_mut();
    let heap_obj = heap::HeapObject::with_data(buffer);
    let heap_idx = state.heap.insert(heap_obj);

    unsafe { std::mem::transmute::<heap::HeapIdx, u64>(heap_idx) as i64 }
}

fn box_load_host(caller: Caller<'_, RuntimeState>, dst_ptr: i32, obj: i64, size_words: i32) {
    log::trace!("box_load_host: dst_ptr={dst_ptr}, obj={obj}, size_words={size_words}");

    // Same as heap_load
    heap_load_host(caller, dst_ptr, obj, size_words);
}

fn box_clone_host(mut caller: Caller<'_, RuntimeState>, obj: i64) {
    log::trace!("box_clone_host: obj={obj}");

    // Same as heap_retain
    let state = caller.data_mut();
    let heap_idx: heap::HeapIdx = unsafe { std::mem::transmute::<u64, heap::HeapIdx>(obj as u64) };
    heap::heap_retain(&mut state.heap, heap_idx);
}

fn box_release_host(mut caller: Caller<'_, RuntimeState>, obj: i64) {
    log::trace!("box_release_host: obj={obj}");

    // Same as heap_release
    let state = caller.data_mut();
    let heap_idx: heap::HeapIdx = unsafe { std::mem::transmute::<u64, heap::HeapIdx>(obj as u64) };
    heap::heap_release(&mut state.heap, heap_idx);
}

fn box_store_host(caller: Caller<'_, RuntimeState>, obj: i64, src_ptr: i32, size_words: i32) {
    log::trace!("box_store_host: obj={obj}, src_ptr={src_ptr}, size_words={size_words}");

    // Same as heap_store
    heap_store_host(caller, obj, src_ptr, size_words);
}

// UserSum operations

/// Copy a user-defined sum type value within linear memory.
/// The wasmgen emits `CloneUserSum` for values that may contain
/// heap-allocated inner payloads  Efor now we perform a shallow
/// byte-level copy (deep clone of boxed payloads is not yet handled).
fn usersum_clone_host(mut caller: Caller<'_, RuntimeState>, dst_ptr: i32, src_ptr: i32, size: i32) {
    log::trace!("usersum_clone_host: dst_ptr={dst_ptr}, src_ptr={src_ptr}, size={size}");

    if size <= 0 || src_ptr == dst_ptr {
        return;
    }

    let byte_len = size as usize * std::mem::size_of::<Word>();
    let memory = caller.data().memory.expect("Memory not initialized");

    // Read source bytes, then write to destination
    let mut buf = vec![0u8; byte_len];
    memory
        .read(&caller, src_ptr as usize, &mut buf)
        .expect("usersum_clone: failed to read source");
    memory
        .write(&mut caller, dst_ptr as usize, &buf)
        .expect("usersum_clone: failed to write destination");
}

/// Release a user-defined sum type value.
/// Currently a no-op because the wasmgen emits placeholder arguments
/// (ptr=0, tag=0). Once the wasmgen passes real pointers and the
/// tag-based dispatch is implemented, this function should inspect
/// the tag, and if the active variant holds a heap-allocated payload,
/// call the corresponding heap_release.
fn usersum_release_host(_caller: Caller<'_, RuntimeState>, ptr: i32, tag: i32, size: i32) {
    log::trace!("usersum_release_host: ptr={ptr}, tag={tag}, size={size}");
}

// Closure operations
//
// NOTE: closure_make, closure_close, and closure_call are registered as WASM
// imports but are intentionally **never emitted** by the current wasmgen.
// The WASM backend handles closures via `call_indirect` on the function table,
// bypassing these host functions entirely. They remain as imports only for
// forward compatibility.

fn closure_make_host(
    _caller: Caller<'_, RuntimeState>,
    fn_idx: i64,
    captured_ptr: i32,
    captured_size: i32,
) -> i64 {
    log::trace!(
        "closure_make_host: fn_idx={fn_idx}, captured_ptr={captured_ptr}, captured_size={captured_size}"
    );
    0
}

fn closure_close_host(_caller: Caller<'_, RuntimeState>, closure: i64) {
    log::trace!("closure_close_host: closure={closure}");
}

fn closure_call_host(
    _caller: Caller<'_, RuntimeState>,
    closure: i64,
    args_ptr: i32,
    args_size: i32,
    dst_ptr: i32,
    dst_size: i32,
) {
    log::trace!(
        "closure_call_host: closure={closure}, args_ptr={args_ptr}, args_size={args_size}, dst_ptr={dst_ptr}, dst_size={dst_size}"
    );
}

// Closure state stack operations
// These mirror the native VM's states_stack.push/pop pattern.
// When a closure is called, its state storage is activated;
// when the call returns, the previous state context is restored.

/// Push a closure's state onto the state stack, making it the active state context.
/// Called before a closure call (CallCls/CallIndirect).
/// `closure_addr` is the closure's address in WASM linear memory (used as a unique key).
/// `state_size` is the number of words needed for this closure's state storage
/// (computed from state_skeleton.total_size() at compile time).
fn closure_state_push_host(
    mut caller: Caller<'_, RuntimeState>,
    closure_addr: i64,
    state_size: i64,
) {
    log::trace!("closure_state_push_host: closure_addr={closure_addr}, state_size={state_size}");
    let state = caller.data_mut();

    // Push closure address onto the state stack
    state.state_stack.push(closure_addr);
    // Lazily allocate state storage for this closure if it doesn't exist yet
    state
        .closure_states
        .entry(closure_addr)
        .or_insert_with(|| StateStorage::with_size(state_size as usize));
}

/// Pop the closure state stack, restoring the previous state context.
/// Called after a closure call returns.
fn closure_state_pop_host(mut caller: Caller<'_, RuntimeState>) {
    log::trace!("closure_state_pop_host");
    let state = caller.data_mut();
    // Reset pos of the outgoing closure's state for next call
    if let Some(&closure_addr) = state.state_stack.last()
        && let Some(cls_state) = state.closure_states.get_mut(&closure_addr)
    {
        cls_state.pos = 0;
    }
    state.state_stack.pop();
}

// State operations

/// Push the state position cursor forward by `offset` words.
/// Mirrors the native VM's `PushStatePos` instruction.
fn state_push_host(mut caller: Caller<'_, RuntimeState>, offset: i64) {
    log::trace!("state_push_host: offset={offset}");
    let state = caller.data_mut();
    let current = state.get_current_state();
    current.pos += offset as usize;
}

/// Pop the state position cursor back by `offset` words.
/// Mirrors the native VM's `PopStatePos` instruction.
fn state_pop_host(mut caller: Caller<'_, RuntimeState>, offset: i64) {
    log::trace!("state_pop_host: offset={offset}");
    let state = caller.data_mut();
    let current = state.get_current_state();
    current.pos -= offset as usize;
}

/// Read state data from the current state storage and write to WASM linear memory.
/// Mirrors the native VM's `GetState` instruction.
fn state_get_host(mut caller: Caller<'_, RuntimeState>, dst_ptr: i32, size_words: i32) {
    log::trace!("state_get_host: dst_ptr={dst_ptr}, size_words={size_words}");

    let size = size_words as usize;

    // Read from the active state storage at its current position
    let state_values: Vec<u64> = {
        let state = caller.data_mut();
        let current = state.get_current_state();
        let pos = current.pos;
        let needed = pos + size;
        if needed > current.data.len() {
            // Grow data to accommodate the requested range (initialized to zero)
            current.data.resize(needed, 0);
        }
        current.data[pos..pos + size].to_vec()
    };

    // Write to WASM linear memory
    let memory = caller.data().memory.expect("Memory not initialized");
    let bytes: Vec<u8> = state_values.iter().flat_map(|w| w.to_le_bytes()).collect();

    memory
        .write(&mut caller.as_context_mut(), dst_ptr as usize, &bytes)
        .expect("Failed to write state to WASM memory");
}

/// Write values from WASM linear memory to the current state storage.
/// Mirrors the native VM's `SetState` instruction.
fn state_set_host(mut caller: Caller<'_, RuntimeState>, src_ptr: i32, size_words: i32) {
    log::trace!("state_set_host: src_ptr={src_ptr}, size_words={size_words}");

    let size = size_words as usize;

    // Read from WASM linear memory
    let state_values: Vec<u64> = {
        let memory = caller.data().memory.expect("Memory not initialized");
        let mut bytes = vec![0u8; size * std::mem::size_of::<Word>()];
        memory
            .read(&caller, src_ptr as usize, &mut bytes)
            .expect("Failed to read from WASM memory");
        bytes
            .chunks(std::mem::size_of::<Word>())
            .map(|chunk| {
                let mut word_bytes = [0u8; 8];
                word_bytes.copy_from_slice(chunk);
                u64::from_le_bytes(word_bytes)
            })
            .collect()
    };

    // Write to the active state storage at its current position
    let state = caller.data_mut();
    let current = state.get_current_state();
    let pos = current.pos;
    let needed = pos + size;

    // Grow data if needed
    if needed > current.data.len() {
        current.data.resize(needed, 0);
    }

    current.data[pos..pos + size].copy_from_slice(&state_values);
}

/// Ring buffer delay: reads delayed value from state, writes new input.
/// State layout at current pos: [read_idx, write_idx, data[0..max_len]]
fn state_delay_host(
    mut caller: Caller<'_, RuntimeState>,
    input: f64,
    time: f64,
    max_len: i64,
) -> f64 {
    let state = caller.data_mut();
    let current = state.get_current_state();
    let pos = current.pos;
    let buf_size = max_len as usize;
    let total_needed = pos + 2 + buf_size;

    // Ensure data is large enough
    if total_needed > current.data.len() {
        current.data.resize(total_needed, 0);
    }

    if buf_size == 0 {
        return 0.0;
    }

    let len = buf_size as u64;
    let max_delay = (len - 1) as f64;
    let delay_samples = time.clamp(0.0, max_delay) as u64;
    let write_idx = current.data[pos + 1] % len;
    let read_idx = (write_idx + len - delay_samples) % len;

    // Read delayed value from ring buffer
    let res_bits = current.data[pos + 2 + read_idx as usize];
    let res = f64::from_bits(res_bits);

    // Write input to ring buffer
    current.data[pos + 2 + write_idx as usize] = input.to_bits();

    // Advance write index and keep current read index for introspection
    current.data[pos] = read_idx;
    current.data[pos + 1] = (write_idx + 1) % len;

    res
}

/// One-sample delay (mem): returns previous value, stores new input.
/// State layout at current pos: [value] (1 word)
fn state_mem_host(mut caller: Caller<'_, RuntimeState>, input: f64) -> f64 {
    let state = caller.data_mut();
    let current = state.get_current_state();
    let pos = current.pos;
    let needed = pos + 1;

    if needed > current.data.len() {
        current.data.resize(needed, 0);
    }

    let old_bits = current.data[pos];
    let old_value = f64::from_bits(old_bits);
    current.data[pos] = input.to_bits();

    old_value
}

// Array operations

fn array_alloc_host(mut caller: Caller<'_, RuntimeState>, src_ptr: i64, size: i32) -> i64 {
    log::trace!("array_alloc_host: src_ptr={src_ptr}, size={size}");

    let size_usize = size as usize;

    // Read from linear memory (src_ptr is i64 but used as memory offset)
    let mut buffer = vec![0u64; size_usize];
    let memory = caller.data().memory.expect("Memory not initialized");
    let offset = src_ptr as usize;
    let bytes = unsafe {
        std::slice::from_raw_parts_mut(
            buffer.as_mut_ptr() as *mut u8,
            size_usize * std::mem::size_of::<Word>(),
        )
    };
    memory
        .read(&caller, offset, bytes)
        .expect("Failed to read from WASM memory");

    // Generate unique array ID
    let state = caller.data_mut();
    let array_id = (state.arrays.len() + 1) as Word;
    state.arrays.insert(array_id, buffer);

    array_id as i64
}

fn array_get_elem_host(
    mut caller: Caller<'_, RuntimeState>,
    dst_ptr: i32,
    array: i64,
    index: i64,
    elem_size: i32,
) {
    log::trace!(
        "array_get_elem_host: dst_ptr={dst_ptr}, array={array}, index={index}, elem_size={elem_size}"
    );

    let idx = index as usize;
    let elem_words = elem_size as usize;

    // Get element data from array storage (multi-word aware)
    let data = {
        let state = caller.data();
        let array_data = state
            .arrays
            .get(&(array as Word))
            .expect("array_get_elem: invalid array ID");
        let base = idx * elem_words;
        if base + elem_words > array_data.len() {
            panic!(
                "array_get_elem: index {} out of bounds (base={}, elem_words={}, len={})",
                idx,
                base,
                elem_words,
                array_data.len()
            );
        }
        array_data[base..base + elem_words].to_vec()
    };

    // Write element data to linear memory at dst_ptr
    let memory = caller.data().memory.expect("Memory not initialized");
    let bytes = unsafe {
        std::slice::from_raw_parts(
            data.as_ptr() as *const u8,
            elem_words * std::mem::size_of::<Word>(),
        )
    };
    memory
        .write(&mut caller, dst_ptr as usize, bytes)
        .expect("Failed to write to WASM memory");
}

fn array_set_elem_host(
    mut caller: Caller<'_, RuntimeState>,
    array: i64,
    index: i64,
    src_ptr: i32,
    elem_size: i32,
) {
    log::trace!(
        "array_set_elem_host: array={array}, index={index}, src_ptr={src_ptr}, elem_size={elem_size}"
    );

    let idx = index as usize;
    let elem_words = elem_size as usize;

    // Read element from linear memory
    let mut buffer = vec![0u64; elem_words];
    let memory = caller.data().memory.expect("Memory not initialized");
    let bytes = unsafe {
        std::slice::from_raw_parts_mut(
            buffer.as_mut_ptr() as *mut u8,
            elem_words * std::mem::size_of::<Word>(),
        )
    };
    memory
        .read(&caller, src_ptr as usize, bytes)
        .expect("Failed to read from WASM memory");

    // Write to array storage (multi-word aware)
    let state = caller.data_mut();
    let array_data = state
        .arrays
        .get_mut(&(array as Word))
        .expect("array_set_elem: invalid array ID");
    let base = idx * elem_words;
    if base + elem_words > array_data.len() {
        panic!("array_set_elem: index out of bounds");
    }
    array_data[base..base + elem_words].copy_from_slice(&buffer);
}

// Runtime globals

fn runtime_get_now_host(caller: Caller<'_, RuntimeState>) -> f64 {
    log::trace!("runtime_get_now_host");
    let state = caller.data();
    state.current_time as f64
}

fn runtime_get_samplerate_host(caller: Caller<'_, RuntimeState>) -> f64 {
    log::trace!("runtime_get_samplerate_host");
    let state = caller.data();
    state.sample_rate
}

// Builtin function host implementations

fn builtin_probeln_host(_caller: Caller<'_, RuntimeState>, x: f64) -> f64 {
    println!("{x}");
    x
}

fn builtin_probe_host(_caller: Caller<'_, RuntimeState>, x: f64) -> f64 {
    print!("{x}");
    x
}

fn builtin_length_array_host(caller: Caller<'_, RuntimeState>, array: i64) -> f64 {
    log::trace!("builtin_length_array_host: array={array}");
    let state = caller.data();
    let array_data = state
        .arrays
        .get(&(array as Word))
        .expect("length_array: invalid array ID");
    array_data.len() as f64
}

fn builtin_split_head_arity_host(
    mut caller: Caller<'_, RuntimeState>,
    array: i64,
    dst_ptr: i32,
    elem_words: usize,
) {
    log::trace!("builtin_split_head$arity{elem_words}: array={array}, dst_ptr={dst_ptr}");

    let (head_words, rest_data) = {
        let state = caller.data();
        let array_data = state
            .arrays
            .get(&(array as Word))
            .expect("split_head: invalid array ID");
        debug_assert!(
            array_data.len() >= elem_words,
            "split_head$arity{elem_words}: array shorter than one element"
        );
        assert!(
            array_data.len() % elem_words == 0,
            "split_head$arity{}: array length {} is not divisible by elem_words",
            elem_words,
            array_data.len()
        );
        let head_words = array_data[..elem_words].to_vec();
        let rest_data = array_data[elem_words..].to_vec();
        (head_words, rest_data)
    };

    let rest_id = {
        let state = caller.data_mut();
        let rest_id = (state.arrays.len() + 1) as Word;
        state.arrays.insert(rest_id, rest_data);
        rest_id
    };

    let memory = caller.data().memory.expect("Memory not initialized");
    head_words.iter().enumerate().for_each(|(idx, word)| {
        memory
            .write(
                &mut caller,
                (dst_ptr as usize) + idx * std::mem::size_of::<Word>(),
                &word.to_le_bytes(),
            )
            .expect("split_head: failed to write head word");
    });
    memory
        .write(
            &mut caller,
            (dst_ptr as usize) + elem_words * std::mem::size_of::<Word>(),
            &(rest_id as i64).to_le_bytes(),
        )
        .expect("split_head: failed to write rest array handle");
}

fn builtin_split_tail_arity_host(
    mut caller: Caller<'_, RuntimeState>,
    array: i64,
    dst_ptr: i32,
    elem_words: usize,
) {
    log::trace!("builtin_split_tail$arity{elem_words}: array={array}, dst_ptr={dst_ptr}");

    let (tail_words, rest_data) = {
        let state = caller.data();
        let array_data = state
            .arrays
            .get(&(array as Word))
            .expect("split_tail: invalid array ID");
        debug_assert!(
            array_data.len() >= elem_words,
            "split_tail$arity{elem_words}: array shorter than one element"
        );
        debug_assert!(
            array_data.len() % elem_words == 0,
            "split_tail$arity{}: array length {} is not divisible by elem_words",
            elem_words,
            array_data.len()
        );
        let tail_start = array_data.len() - elem_words;
        let tail_words = array_data[tail_start..].to_vec();
        let rest_data = array_data[..tail_start].to_vec();
        (tail_words, rest_data)
    };

    let rest_id = {
        let state = caller.data_mut();
        let rest_id = (state.arrays.len() + 1) as Word;
        state.arrays.insert(rest_id, rest_data);
        rest_id
    };

    let memory = caller.data().memory.expect("Memory not initialized");
    memory
        .write(
            &mut caller,
            dst_ptr as usize,
            &(rest_id as i64).to_le_bytes(),
        )
        .expect("split_tail: failed to write rest array handle");
    tail_words.iter().enumerate().for_each(|(idx, word)| {
        memory
            .write(
                &mut caller,
                (dst_ptr as usize) + (idx + 1) * std::mem::size_of::<Word>(),
                &word.to_le_bytes(),
            )
            .expect("split_tail: failed to write tail word");
    });
}

fn builtin_prepend_arity_host(
    mut caller: Caller<'_, RuntimeState>,
    params: &[Val],
    results: &mut [Val],
    elem_words: usize,
) {
    let expected_params = 2;
    if params.len() != expected_params {
        panic!(
            "prepend$arity{} expected {} params, got {}",
            elem_words,
            expected_params,
            params.len()
        );
    }

    let elem_or_ptr = match params[0] {
        Val::I64(i) => i as u64,
        Val::F64(bits) => bits,
        ref other => {
            panic!("prepend$arity{elem_words}: expected first arg as I64/F64, got {other:?}")
        }
    };

    let elem_bits = if elem_words == 1 {
        vec![elem_or_ptr]
    } else {
        let mut words = vec![0u64; elem_words];
        let memory = caller.data().memory.expect("Memory not initialized");
        let src_ptr = elem_or_ptr as usize;
        let bytes = unsafe {
            std::slice::from_raw_parts_mut(
                words.as_mut_ptr() as *mut u8,
                elem_words * std::mem::size_of::<Word>(),
            )
        };
        memory
            .read(&caller, src_ptr, bytes)
            .expect("prepend$arityN: failed to read element words from memory");
        words
    };

    let array_handle = match params[1] {
        Val::I64(i) => i as Word,
        Val::F64(bits) => i64::from_le_bytes(bits.to_le_bytes()) as Word,
        ref other => {
            panic!("prepend$arity{elem_words}: expected array handle as I64/F64, got {other:?}")
        }
    };

    let mut new_array = elem_bits;
    {
        let state = caller.data();
        let old = state
            .arrays
            .get(&array_handle)
            .unwrap_or_else(|| panic!("prepend: invalid array ID {array_handle}"));
        if old.len() % elem_words != 0 {
            panic!(
                "prepend$arity{}: array length {} is not divisible by elem_words",
                elem_words,
                old.len()
            );
        }
        new_array.extend_from_slice(old);
    }

    let new_id = {
        let state = caller.data_mut();
        let new_id = (state.arrays.len() + 1) as Word;
        state.arrays.insert(new_id, new_array);
        new_id
    };

    if let Some(slot) = results.get_mut(0) {
        *slot = Val::I64(new_id as i64);
    }
}

fn builtin_split_head_host(caller: Caller<'_, RuntimeState>, array: i64, dst_ptr: i32) {
    builtin_split_head_arity_host(caller, array, dst_ptr, 1);
}

fn builtin_split_tail_host(caller: Caller<'_, RuntimeState>, array: i64, dst_ptr: i32) {
    builtin_split_tail_arity_host(caller, array, dst_ptr, 1);
}

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

    #[test]
    fn test_wasm_runtime_create() {
        let runtime = WasmRuntime::new(&[], None);
        assert!(runtime.is_ok(), "Should create WASM runtime");
    }

    #[test]
    fn test_wasm_runtime_load_empty_module() {
        let mut runtime = WasmRuntime::new(&[], None).unwrap();

        // Minimal valid WASM module with memory export
        let wasm_bytes = wat::parse_str(
            r#"
            (module
                (memory (export "memory") 1)
            )
            "#,
        )
        .expect("Failed to parse WAT");

        let result = runtime.load_module(&wasm_bytes);
        assert!(result.is_ok(), "Should load minimal WASM module");
    }

    #[test]
    fn test_heap_operations() {
        let mut runtime = WasmRuntime::new(&[], None).unwrap();

        // WASM module that tests heap operations
        let wasm_bytes = wat::parse_str(
            r#"
            (module
                (import "runtime" "heap_alloc" (func $heap_alloc (param i32) (result i64)))
                (import "runtime" "heap_retain" (func $heap_retain (param i64)))
                (import "runtime" "heap_release" (func $heap_release (param i64)))
                (memory (export "memory") 1)
                
                (func (export "test_heap") (result i64)
                    (local $obj i64)
                    ;; Allocate heap object with 2 words
                    i32.const 2
                    call $heap_alloc
                    local.set $obj
                    
                    ;; Retain it
                    local.get $obj
                    call $heap_retain
                    
                    ;; Return the heap index
                    local.get $obj
                )
            )
            "#,
        )
        .expect("Failed to parse WAT");

        let mut module = runtime
            .load_module(&wasm_bytes)
            .expect("Failed to load module");
        let result = module
            .call_function("test_heap", &[])
            .expect("Failed to call function");

        assert_eq!(result.len(), 1, "Should return one value");
        assert_ne!(result[0], 0, "Heap allocation should return non-zero index");
    }

    #[test]
    fn test_array_operations() {
        let mut runtime = WasmRuntime::new(&[], None).unwrap();

        // WASM module that tests array operations
        let wasm_bytes = wat::parse_str(
            r#"
            (module
                (import "runtime" "array_alloc" (func $array_alloc (param i64 i32) (result i64)))
                (import "runtime" "array_get_elem" (func $array_get_elem (param i32 i64 i64 i32)))
                (import "runtime" "array_set_elem" (func $array_set_elem (param i64 i64 i32 i32)))
                (memory (export "memory") 1)
                
                (func (export "test_array") (result i64)
                    (local $arr i64)
                    ;; Store some values in memory
                    i32.const 0
                    i64.const 42
                    i64.store
                    i32.const 8
                    i64.const 99
                    i64.store
                    
                    ;; Allocate array from memory (ptr=0 as i64, size=2)
                    i64.const 0
                    i32.const 2
                    call $array_alloc
                    local.set $arr
                    
                    ;; Get first element into memory at offset 16
                    ;; (dst_ptr=16, array, index=0 as i64, elem_size=1)
                    i32.const 16
                    local.get $arr
                    i64.const 0
                    i32.const 1
                    call $array_get_elem

                    ;; Load the result from memory offset 16
                    i32.const 16
                    i64.load
                )
            )
            "#,
        )
        .expect("Failed to parse WAT");

        let mut module = runtime
            .load_module(&wasm_bytes)
            .expect("Failed to load module");
        let result = module
            .call_function("test_array", &[])
            .expect("Failed to call function");

        assert_eq!(result.len(), 1, "Should return one value");
        assert_eq!(result[0], 42, "Array should contain stored value");
    }

    #[test]
    fn test_math_hyperbolic_imports() {
        let mut runtime = WasmRuntime::new(&[], None).unwrap();

        let wasm_bytes = wat::parse_str(
            r#"
            (module
                (import "math" "sinh" (func $sinh (param f64) (result f64)))
                (import "math" "cosh" (func $cosh (param f64) (result f64)))
                (import "math" "tanh" (func $tanh (param f64) (result f64)))
                (memory (export "memory") 1)

                (func (export "test_sinh") (result f64)
                    f64.const 0.5
                    call $sinh)

                (func (export "test_cosh") (result f64)
                    f64.const 0.5
                    call $cosh)

                (func (export "test_tanh") (result f64)
                    f64.const 0.5
                    call $tanh)
            )
            "#,
        )
        .expect("Failed to parse WAT");

        let mut module = runtime
            .load_module(&wasm_bytes)
            .expect("Failed to load module");

        let sinh_res = module
            .call_function("test_sinh", &[])
            .expect("Failed to call test_sinh");
        let cosh_res = module
            .call_function("test_cosh", &[])
            .expect("Failed to call test_cosh");
        let tanh_res = module
            .call_function("test_tanh", &[])
            .expect("Failed to call test_tanh");

        let sinh_val = f64::from_bits(sinh_res[0]);
        let cosh_val = f64::from_bits(cosh_res[0]);
        let tanh_val = f64::from_bits(tanh_res[0]);

        assert!((sinh_val - 0.5f64.sinh()).abs() < 1e-12);
        assert!((cosh_val - 0.5f64.cosh()).abs() < 1e-12);
        assert!((tanh_val - 0.5f64.tanh()).abs() < 1e-12);
    }

    #[test]
    fn test_load_generated_wasm() {
        use std::path::PathBuf;

        // Use workspace root relative path
        let mut wasm_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
        wasm_path.pop(); // Remove mimium-lang
        wasm_path.pop(); // Remove lib
        wasm_path.pop(); // Remove crates
        wasm_path.push("tmp");
        wasm_path.push("test_integration.wasm");

        // Skip test if the file doesn't exist (e.g., in CI environment)
        if !wasm_path.exists() {
            eprintln!("Skipping test_load_generated_wasm: file not found at {wasm_path:?}");
            return;
        }

        eprintln!("Loading WASM from: {wasm_path:?}");
        let wasm_bytes = std::fs::read(&wasm_path).expect("Failed to read generated WASM file");

        let mut runtime = WasmRuntime::new(&[], None).unwrap();
        let result = runtime.load_module(&wasm_bytes);

        assert!(
            result.is_ok(),
            "Should load generated WASM module: {:?}",
            result.err()
        );

        let mut module = result.unwrap();

        // Try to call the dsp function if it exists
        if let Ok(result) = module.call_function("fn_0", &[]) {
            eprintln!("fn_0 returned: {result:?}");
        }
    }
}