use clusterflux_core::{
    DebugEpoch, DebugParticipant, DebugParticipantKind, DebugRuntimeState, Digest, ProcessId,
    TaskInstanceId, WasmHostCommandRequest, WasmHostCommandResult, WasmHostDebugProbeRequest,
    WasmHostDebugProbeResult, WasmHostSourceSnapshotRequest, WasmHostSourceSnapshotResult,
    WasmHostTaskControlRequest, WasmHostTaskControlResult, WasmHostTaskHandle,
    WasmHostTaskJoinRequest, WasmHostTaskJoinResult, WasmHostTaskStartRequest, WasmHostVfsRequest,
    WasmHostVfsResult, WasmTaskInvocation, WasmTaskResult, MAX_WASM_TASK_ENVELOPE_BYTES,
};
use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Condvar, Mutex};
use std::thread;
use std::time::{Duration, Instant};
use thiserror::Error;
use wasmtime::{
    Caller, Config, DebugEvent, DebugHandler, Engine, Linker, Module, OptLevel, Store,
    StoreContextMut, StoreLimits, StoreLimitsBuilder, UpdateDeadline, WasmBacktrace,
};

mod task_host_linker;
use task_host_linker::{task_host_linker, task_host_stub_linker};

const INACTIVE_EPOCH_DEADLINE_TICKS: u64 = u64::MAX / 2;
const DEFAULT_MAX_WASM_MEMORY_BYTES: usize = 256 * 1024 * 1024;
const MAX_WASM_TABLE_ELEMENTS: usize = 100_000;
const DEFAULT_WASM_FUEL_UNITS_PER_SECOND: u64 = 10_000_000;
const DEFAULT_WASM_FUEL_BURST_SECONDS: u64 = 60;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WasmtimeRuntimeLimits {
    pub fuel_units_per_second: u64,
    pub fuel_burst_seconds: u64,
    pub memory_bytes: usize,
}

impl WasmtimeRuntimeLimits {
    pub fn validate(&self) -> Result<(), WasmTaskError> {
        if self.fuel_units_per_second == 0 || self.fuel_burst_seconds == 0 {
            return Err(WasmTaskError::Runtime(
                "Wasm fuel rate and burst duration must be positive".to_owned(),
            ));
        }
        if self.memory_bytes == 0 {
            return Err(WasmTaskError::Runtime(
                "Wasm memory limit must be positive".to_owned(),
            ));
        }
        self.fuel_capacity().ok_or_else(|| {
            WasmTaskError::Runtime("Wasm fuel burst capacity overflowed u64".to_owned())
        })?;
        Ok(())
    }

    fn fuel_capacity(&self) -> Option<u64> {
        self.fuel_units_per_second
            .checked_mul(self.fuel_burst_seconds)
    }
}

impl Default for WasmtimeRuntimeLimits {
    fn default() -> Self {
        Self {
            fuel_units_per_second: DEFAULT_WASM_FUEL_UNITS_PER_SECOND,
            fuel_burst_seconds: DEFAULT_WASM_FUEL_BURST_SECONDS,
            memory_bytes: DEFAULT_MAX_WASM_MEMORY_BYTES,
        }
    }
}

fn task_store_limits(runtime: &WasmtimeRuntimeLimits) -> StoreLimits {
    StoreLimitsBuilder::new()
        .memory_size(runtime.memory_bytes)
        .table_elements(MAX_WASM_TABLE_ELEMENTS)
        .instances(8)
        .tables(8)
        .memories(8)
        .trap_on_grow_failure(true)
        .build()
}

fn debug_control_trace(message: impl std::fmt::Display) {
    if std::env::var_os("CLUSTERFLUX_DEBUG_CONTROL_TRACE").is_some() {
        eprintln!("clusterflux debug control: {message}");
    }
}

#[derive(Clone)]
pub struct WasmtimeTaskRuntime {
    engine: Engine,
    runtime_limits: WasmtimeRuntimeLimits,
}

pub trait WasmTaskHost {
    fn abort_signal(&self) -> Option<Arc<AtomicBool>> {
        None
    }

    fn debug_control(&self) -> Option<Arc<WasmDebugControl>> {
        None
    }

    fn start_task(
        &mut self,
        request: WasmHostTaskStartRequest,
    ) -> Result<WasmHostTaskHandle, String>;
    fn join_task(
        &mut self,
        request: WasmHostTaskJoinRequest,
    ) -> Result<WasmHostTaskJoinResult, String>;
    fn run_command(
        &mut self,
        request: WasmHostCommandRequest,
    ) -> Result<WasmHostCommandResult, String>;
    fn poll_task_control(
        &mut self,
        request: WasmHostTaskControlRequest,
    ) -> Result<WasmHostTaskControlResult, String>;
    fn debug_probe(
        &mut self,
        request: WasmHostDebugProbeRequest,
    ) -> Result<WasmHostDebugProbeResult, String> {
        request.validate()?;
        Ok(WasmHostDebugProbeResult {
            abi_version: clusterflux_core::WASM_TASK_ABI_VERSION,
            breakpoint_matched: false,
            debug_epoch: None,
        })
    }
    fn vfs_operation(&mut self, request: WasmHostVfsRequest) -> Result<WasmHostVfsResult, String>;
    fn snapshot_source(
        &mut self,
        request: WasmHostSourceSnapshotRequest,
    ) -> Result<WasmHostSourceSnapshotResult, String>;
}

#[derive(Debug, Default)]
struct WasmDebugControlState {
    requested_epoch: Option<u64>,
    frozen_epoch: Option<u64>,
    resumed_through_epoch: u64,
    execution_armed: bool,
    quiescent_host_boundary_depth: usize,
    frozen_at_host_boundary: bool,
    stack_frames: Vec<String>,
}

#[derive(Debug, Default)]
pub struct WasmDebugControl {
    state: Mutex<WasmDebugControlState>,
    changed: Condvar,
}

impl WasmDebugControl {
    pub fn request_freeze(&self, epoch: u64) {
        let mut state = self.state.lock().expect("debug control lock poisoned");
        if state
            .requested_epoch
            .is_none_or(|requested| epoch >= requested)
        {
            state.requested_epoch = Some(epoch);
            if (!state.execution_armed || state.quiescent_host_boundary_depth > 0)
                && state.resumed_through_epoch < epoch
            {
                state.frozen_epoch = Some(epoch);
                state.frozen_at_host_boundary = true;
            }
            self.changed.notify_all();
        }
    }

    pub fn request_resume(&self, epoch: u64) {
        let mut state = self.state.lock().expect("debug control lock poisoned");
        state.resumed_through_epoch = state.resumed_through_epoch.max(epoch);
        if state.frozen_epoch == Some(epoch) && state.frozen_at_host_boundary {
            state.frozen_epoch = None;
            state.frozen_at_host_boundary = false;
        }
        self.changed.notify_all();
    }

    pub fn requested_epoch(&self) -> Option<u64> {
        self.state
            .lock()
            .expect("debug control lock poisoned")
            .requested_epoch
    }

    pub fn frozen_epoch(&self) -> Option<u64> {
        self.state
            .lock()
            .expect("debug control lock poisoned")
            .frozen_epoch
    }

    pub fn resume_requested(&self, epoch: u64) -> bool {
        self.state
            .lock()
            .expect("debug control lock poisoned")
            .resumed_through_epoch
            >= epoch
    }

    pub fn mark_frozen(&self, epoch: u64) {
        let mut state = self.state.lock().expect("debug control lock poisoned");
        state.frozen_epoch = Some(epoch);
        state.frozen_at_host_boundary = false;
        self.changed.notify_all();
    }

    pub fn mark_running(&self, epoch: u64) {
        let mut state = self.state.lock().expect("debug control lock poisoned");
        if state.frozen_epoch == Some(epoch) {
            state.frozen_epoch = None;
            state.frozen_at_host_boundary = false;
        }
        self.changed.notify_all();
    }

    pub fn wait_until_frozen(&self, epoch: u64, timeout: Duration) -> bool {
        let state = self.state.lock().expect("debug control lock poisoned");
        let (state, _) = self
            .changed
            .wait_timeout_while(state, timeout, |state| state.frozen_epoch != Some(epoch))
            .expect("debug control lock poisoned while waiting for freeze");
        state.frozen_epoch == Some(epoch)
    }

    pub fn wait_until_running(&self, epoch: u64, timeout: Duration) -> bool {
        let state = self.state.lock().expect("debug control lock poisoned");
        let (state, _) = self
            .changed
            .wait_timeout_while(state, timeout, |state| state.frozen_epoch == Some(epoch))
            .expect("debug control lock poisoned while waiting for resume");
        state.frozen_epoch != Some(epoch)
    }

    pub fn stack_frames(&self) -> Vec<String> {
        self.state
            .lock()
            .expect("debug control lock poisoned")
            .stack_frames
            .clone()
    }

    fn record_stack_frames(&self, stack_frames: Vec<String>) {
        self.state
            .lock()
            .expect("debug control lock poisoned")
            .stack_frames = stack_frames;
    }

    fn arm_execution(&self, abort: &AtomicBool) {
        let mut state = self.state.lock().expect("debug control lock poisoned");
        state.execution_armed = true;
        while state.frozen_at_host_boundary
            && state.frozen_epoch.is_some()
            && !abort.load(Ordering::Acquire)
        {
            state = self
                .changed
                .wait_timeout(state, Duration::from_millis(50))
                .expect("debug control lock poisoned at Wasm startup safepoint")
                .0;
        }
    }

    fn pause_at_requested_epoch(&self, abort: &AtomicBool) {
        let mut state = self.state.lock().expect("debug control lock poisoned");
        let Some(epoch) = state.requested_epoch else {
            return;
        };
        if state.resumed_through_epoch >= epoch {
            return;
        }
        debug_control_trace(format_args!("Wasm epoch callback freezing epoch {epoch}"));
        state.frozen_epoch = Some(epoch);
        state.frozen_at_host_boundary = false;
        self.changed.notify_all();
        while state.resumed_through_epoch < epoch && !abort.load(Ordering::Acquire) {
            state = self
                .changed
                .wait_timeout(state, Duration::from_millis(50))
                .expect("debug control lock poisoned while Wasm was frozen")
                .0;
        }
        state.frozen_epoch = None;
        debug_control_trace(format_args!("Wasm epoch callback resumed epoch {epoch}"));
        self.changed.notify_all();
    }

    fn enter_quiescent_host_boundary(&self, abort: Option<&AtomicBool>) {
        let mut state = self.state.lock().expect("debug control lock poisoned");
        state.quiescent_host_boundary_depth += 1;
        if let Some(epoch) = state.requested_epoch {
            if state.resumed_through_epoch < epoch {
                state.frozen_epoch = Some(epoch);
                state.frozen_at_host_boundary = true;
                self.changed.notify_all();
            }
        }
        while state.frozen_at_host_boundary
            && state.frozen_epoch.is_some()
            && !abort.is_some_and(|abort| abort.load(Ordering::Acquire))
        {
            state = self
                .changed
                .wait_timeout(state, Duration::from_millis(50))
                .expect("debug control lock poisoned entering host-call safepoint")
                .0;
        }
    }

    fn leave_quiescent_host_boundary(&self, abort: Option<&AtomicBool>) {
        let mut state = self.state.lock().expect("debug control lock poisoned");
        while state.frozen_at_host_boundary
            && state.frozen_epoch.is_some()
            && !abort.is_some_and(|abort| abort.load(Ordering::Acquire))
        {
            state = self
                .changed
                .wait_timeout(state, Duration::from_millis(50))
                .expect("debug control lock poisoned at host-call safepoint")
                .0;
        }
        state.quiescent_host_boundary_depth = state.quiescent_host_boundary_depth.saturating_sub(1);
        if state.quiescent_host_boundary_depth == 0 && state.frozen_at_host_boundary {
            state.frozen_epoch = None;
            state.frozen_at_host_boundary = false;
            self.changed.notify_all();
        }
    }
}

struct WasmtimeTaskHostState {
    host: Box<dyn WasmTaskHost>,
    fatal_host_error: Option<String>,
    limits: StoreLimits,
    fuel_budget: FuelTokenBucket,
}

struct BasicStoreState {
    limits: StoreLimits,
}

struct FuelTokenBucket {
    fuel_units_per_second: u64,
    capacity: u64,
    last_refill: Instant,
    fractional_fuel_numerator: u128,
}

impl FuelTokenBucket {
    fn new(limits: &WasmtimeRuntimeLimits) -> Self {
        Self {
            fuel_units_per_second: limits.fuel_units_per_second,
            capacity: limits
                .fuel_capacity()
                .expect("validated Wasm fuel capacity"),
            last_refill: Instant::now(),
            fractional_fuel_numerator: 0,
        }
    }

    fn refill(&mut self, current_fuel: u64) -> u64 {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_refill);
        self.last_refill = now;
        self.refill_after(current_fuel, elapsed)
    }

    fn refill_after(&mut self, current_fuel: u64, elapsed: Duration) -> u64 {
        let earned_numerator = elapsed
            .as_nanos()
            .saturating_mul(self.fuel_units_per_second as u128)
            .saturating_add(self.fractional_fuel_numerator);
        let refill = earned_numerator
            .checked_div(1_000_000_000)
            .unwrap_or(0)
            .min(u64::MAX as u128) as u64;
        let refilled = current_fuel.saturating_add(refill).min(self.capacity);
        self.fractional_fuel_numerator = if refilled == self.capacity {
            0
        } else {
            earned_numerator % 1_000_000_000
        };
        refilled
    }
}

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

    #[test]
    fn frequent_refills_preserve_fractional_credit() {
        let limits = WasmtimeRuntimeLimits {
            fuel_units_per_second: 10,
            fuel_burst_seconds: 60,
            memory_bytes: 1024,
        };
        let mut bucket = FuelTokenBucket::new(&limits);
        let mut fuel = 0;
        for _ in 0..4 {
            fuel = bucket.refill_after(fuel, Duration::from_millis(25));
        }
        assert_eq!(fuel, 1);
        assert_eq!(bucket.fractional_fuel_numerator, 0);
    }
}

impl WasmtimeTaskHostState {
    pub(crate) fn refill_fuel_after_host_call(&mut self, current_fuel: u64) -> u64 {
        self.fuel_budget.refill(current_fuel)
    }
}

struct EpochControlGuard {
    stop: Arc<AtomicBool>,
    watcher: Option<thread::JoinHandle<()>>,
}

impl EpochControlGuard {
    fn arm<T: 'static>(
        engine: &Engine,
        store: &mut Store<T>,
        abort: Arc<AtomicBool>,
        debug: Option<Arc<WasmDebugControl>>,
    ) -> Self {
        debug_control_trace(format_args!(
            "arming epoch control (debug={:?})",
            debug.as_ref().map(Arc::as_ptr)
        ));
        let callback_abort = Arc::clone(&abort);
        let callback_debug = debug.clone();
        store.epoch_deadline_callback(move |_store| {
            if callback_abort.load(Ordering::Acquire) {
                return Ok(UpdateDeadline::Interrupt);
            }
            if let Some(debug) = &callback_debug {
                debug.pause_at_requested_epoch(&callback_abort);
                if callback_abort.load(Ordering::Acquire) {
                    return Ok(UpdateDeadline::Interrupt);
                }
            }
            // Epochs are engine-wide. Another task may increment the shared
            // engine's epoch to control its own store; keep this store running
            // unless this store's own control signal requires action.
            Ok(UpdateDeadline::Continue(1))
        });
        store.set_epoch_deadline(1);
        if let Some(debug) = &debug {
            debug.arm_execution(&abort);
        }
        let engine = engine.clone();
        let stop = Arc::new(AtomicBool::new(false));
        let watcher_stop = Arc::clone(&stop);
        let watcher = thread::spawn(move || {
            debug_control_trace("epoch control watcher started");
            let mut triggered_debug_epoch = None;
            while !watcher_stop.load(Ordering::Acquire) {
                if abort.load(Ordering::Acquire) {
                    engine.increment_epoch();
                    return;
                }
                if let Some(epoch) = debug.as_ref().and_then(|debug| debug.requested_epoch()) {
                    if triggered_debug_epoch != Some(epoch) {
                        triggered_debug_epoch = Some(epoch);
                        debug_control_trace(format_args!(
                            "incrementing engine epoch for debug epoch {epoch}"
                        ));
                        engine.increment_epoch();
                    }
                }
                thread::sleep(Duration::from_millis(10));
            }
        });
        Self {
            stop,
            watcher: Some(watcher),
        }
    }
}

fn defer_epoch_interruption<T>(store: &mut Store<T>) {
    // Stores default to an already-expired epoch deadline. ABI setup and stores
    // without an abort signal must explicitly opt out of interruption, especially
    // after a previous task has incremented the shared engine epoch.
    store.set_epoch_deadline(INACTIVE_EPOCH_DEADLINE_TICKS);
}

impl Drop for EpochControlGuard {
    fn drop(&mut self) {
        self.stop.store(true, Ordering::Release);
        if let Some(watcher) = self.watcher.take() {
            let _ = watcher.join();
        }
    }
}

fn abort_error(
    store: &Store<WasmtimeTaskHostState>,
    abort: Option<&Arc<AtomicBool>>,
    error: wasmtime::Error,
) -> WasmTaskError {
    if abort.is_some_and(|signal| signal.load(Ordering::Acquire)) {
        return WasmTaskError::HostControl(
            "task execution cancelled: coordinator requested process abort".to_owned(),
        );
    }
    store
        .data()
        .fatal_host_error
        .clone()
        .map(WasmTaskError::HostControl)
        .unwrap_or_else(|| wasmtime_error(error))
}

#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum WasmTaskError {
    #[error("wasmtime task failed: {0}")]
    Runtime(String),
    #[error("Wasm task ABI failed: {0}")]
    TaskAbi(String),
    #[error("{0}")]
    HostControl(String),
    #[error(
        "bundle digest mismatch before wasmtime execution: expected {expected}, actual {actual}"
    )]
    BundleDigestMismatch { expected: Digest, actual: Digest },
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WasmtimeDebugProbe {
    pub task: TaskInstanceId,
    pub frozen_state: DebugRuntimeState,
    pub resumed_state: DebugRuntimeState,
    pub result: i32,
    pub stack_frames: Vec<String>,
    pub local_values: Vec<(String, String)>,
    pub wasm_function: Option<String>,
    pub wasm_pc: Option<u32>,
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
struct WasmtimeFrameSnapshot {
    stack_frames: Vec<String>,
    local_values: Vec<(String, String)>,
    wasm_function: Option<String>,
    wasm_pc: Option<u32>,
}

#[derive(Debug)]
struct WasmtimeDebugState {
    snapshot: Option<WasmtimeFrameSnapshot>,
    limits: StoreLimits,
}

impl Default for WasmtimeDebugState {
    fn default() -> Self {
        Self {
            snapshot: None,
            limits: task_store_limits(&WasmtimeRuntimeLimits::default()),
        }
    }
}

#[derive(Clone, Debug)]
struct WasmtimeLocalSnapshotHandler {
    export: String,
}

impl DebugHandler for WasmtimeLocalSnapshotHandler {
    type Data = WasmtimeDebugState;

    async fn handle(&self, mut store: StoreContextMut<'_, Self::Data>, _event: DebugEvent<'_>) {
        if store.data().snapshot.is_some() {
            if let Some(mut edit) = store.edit_breakpoints() {
                let _ = edit.single_step(false);
            }
            return;
        }

        let mut snapshot = WasmtimeFrameSnapshot {
            stack_frames: vec![format!("{}::wasm_export", self.export)],
            ..WasmtimeFrameSnapshot::default()
        };
        // The debugger only presents the top Wasm frame here. Bounding the iterator is
        // also essential for modules with linked host imports, whose exit-frame chain
        // may include runtime trampolines that are not user-visible frames.
        if let Some(frame) = store
            .debug_exit_frames()
            .take(1)
            .collect::<Vec<_>>()
            .into_iter()
            .next()
        {
            if let Ok(Some((function, pc))) = frame.wasm_function_index_and_pc(&mut store) {
                snapshot.wasm_function = Some(format!("{function:?}"));
                snapshot.wasm_pc = Some(pc);
            }

            if let Ok(count) = frame.num_locals(&mut store) {
                for index in 0..count.min(16) {
                    let value = match frame.local(&mut store, index) {
                        Ok(value) => format!("{value:?}"),
                        Err(err) => format!("<error: {err:#}>"),
                    };
                    snapshot
                        .local_values
                        .push((format!("wasm_local_{index}"), value));
                }
            }
        }

        if let Some(function) = &snapshot.wasm_function {
            snapshot.stack_frames = vec![format!("{} / {function}", self.export)];
        }
        store.data_mut().snapshot = Some(snapshot);
        if let Some(mut edit) = store.edit_breakpoints() {
            let _ = edit.single_step(false);
        }
    }
}

impl WasmtimeTaskRuntime {
    pub fn new() -> Result<Self, WasmTaskError> {
        Self::new_with_limits(WasmtimeRuntimeLimits::default())
    }

    pub fn new_with_limits(runtime_limits: WasmtimeRuntimeLimits) -> Result<Self, WasmTaskError> {
        runtime_limits.validate()?;
        let mut config = Config::new();
        config.wasm_backtrace_details(wasmtime::WasmBacktraceDetails::Enable);
        config.epoch_interruption(true);
        config.consume_fuel(true);
        let engine = Engine::new(&config).map_err(wasmtime_error)?;
        Ok(Self {
            engine,
            runtime_limits,
        })
    }

    fn debug_engine() -> Result<Engine, WasmTaskError> {
        let mut config = Config::new();
        config.debug_info(true);
        config.guest_debug(true);
        config.generate_address_map(true);
        config.cranelift_opt_level(OptLevel::None);
        Engine::new(&config).map_err(wasmtime_error)
    }

    pub fn run_i32_export(
        &self,
        wasm_or_wat: impl AsRef<[u8]>,
        export: &str,
        arg: i32,
    ) -> Result<i32, WasmTaskError> {
        self.run_i32_export_bytes(wasm_or_wat.as_ref(), export, arg)
    }

    pub fn run_i32_export_verified(
        &self,
        wasm_or_wat: impl AsRef<[u8]>,
        expected_bundle_digest: &Digest,
        export: &str,
        arg: i32,
    ) -> Result<i32, WasmTaskError> {
        let wasm_or_wat = Self::verified_module_bytes(wasm_or_wat, expected_bundle_digest)?;
        self.run_i32_export_bytes(&wasm_or_wat, export, arg)
    }

    pub fn run_i32_export_verified_with_task_host(
        &self,
        wasm_or_wat: impl AsRef<[u8]>,
        expected_bundle_digest: &Digest,
        export: &str,
        arg: i32,
        host: Box<dyn WasmTaskHost>,
    ) -> Result<i32, WasmTaskError> {
        let module_bytes = Self::verified_module_bytes(wasm_or_wat, expected_bundle_digest)?;
        let module = Module::new(&self.engine, &module_bytes).map_err(wasmtime_error)?;
        let linker = task_host_linker(&self.engine)?;
        let abort_signal = host.abort_signal();
        let debug_control = host.debug_control();
        let mut store = Store::new(
            &self.engine,
            WasmtimeTaskHostState {
                host,
                fatal_host_error: None,
                limits: task_store_limits(&self.runtime_limits),
                fuel_budget: FuelTokenBucket::new(&self.runtime_limits),
            },
        );
        store.limiter(|state| &mut state.limits);
        store
            .set_fuel(
                self.runtime_limits
                    .fuel_capacity()
                    .expect("validated fuel capacity"),
            )
            .map_err(wasmtime_error)?;
        defer_epoch_interruption(&mut store);
        let instance = linker
            .instantiate(&mut store, &module)
            .map_err(wasmtime_error)?;
        let function = instance
            .get_typed_func::<i32, i32>(&mut store, export)
            .map_err(wasmtime_error)?;
        let _abort_guard = abort_signal.as_ref().map(|signal| {
            EpochControlGuard::arm(&self.engine, &mut store, Arc::clone(signal), debug_control)
        });
        match function.call(&mut store, arg) {
            Ok(result) => Ok(result),
            Err(error) => Err(abort_error(&store, abort_signal.as_ref(), error)),
        }
    }

    pub fn run_task_export_verified(
        &self,
        wasm_or_wat: impl AsRef<[u8]>,
        expected_bundle_digest: &Digest,
        export: &str,
        invocation: &WasmTaskInvocation,
    ) -> Result<WasmTaskResult, WasmTaskError> {
        invocation.validate().map_err(WasmTaskError::TaskAbi)?;
        let module_bytes = Self::verified_module_bytes(wasm_or_wat, expected_bundle_digest)?;
        let encoded = serde_json::to_vec(invocation)
            .map_err(|error| WasmTaskError::TaskAbi(error.to_string()))?;
        let module = Module::new(&self.engine, &module_bytes).map_err(wasmtime_error)?;
        let linker = task_host_stub_linker(&self.engine)?;
        let mut store = Store::new(
            &self.engine,
            BasicStoreState {
                limits: task_store_limits(&self.runtime_limits),
            },
        );
        store.limiter(|state| &mut state.limits);
        store
            .set_fuel(
                self.runtime_limits
                    .fuel_capacity()
                    .expect("validated fuel capacity"),
            )
            .map_err(wasmtime_error)?;
        defer_epoch_interruption(&mut store);
        let instance = linker
            .instantiate(&mut store, &module)
            .map_err(wasmtime_error)?;
        let memory = instance
            .get_memory(&mut store, "memory")
            .ok_or_else(|| WasmTaskError::TaskAbi("guest module exports no memory".to_owned()))?;
        let allocate = instance
            .get_typed_func::<u32, u32>(&mut store, "clusterflux_alloc_v1")
            .map_err(wasmtime_error)?;
        let input_length = u32::try_from(encoded.len())
            .map_err(|_| WasmTaskError::TaskAbi("task invocation is too large".to_owned()))?;
        let input_pointer = allocate
            .call(&mut store, input_length)
            .map_err(wasmtime_error)?;
        debug_control_trace("task ABI input allocation completed");
        if input_pointer == 0 && input_length != 0 {
            return Err(WasmTaskError::TaskAbi(
                "guest refused task invocation allocation".to_owned(),
            ));
        }
        memory
            .write(&mut store, input_pointer as usize, &encoded)
            .map_err(|error| WasmTaskError::TaskAbi(error.to_string()))?;
        let task = instance
            .get_typed_func::<(u32, u32), u64>(&mut store, export)
            .map_err(wasmtime_error)?;
        let packed = task
            .call(&mut store, (input_pointer, input_length))
            .map_err(wasmtime_error)?;
        let result_pointer = packed as u32;
        let result_length = (packed >> 32) as u32;
        if result_length as usize > MAX_WASM_TASK_ENVELOPE_BYTES {
            return Err(WasmTaskError::TaskAbi(format!(
                "guest task result is {result_length} bytes; maximum is {MAX_WASM_TASK_ENVELOPE_BYTES}"
            )));
        }
        if result_pointer == 0 && result_length != 0 {
            return Err(WasmTaskError::TaskAbi(
                "guest returned a null task result pointer".to_owned(),
            ));
        }
        let mut result_bytes = vec![0_u8; result_length as usize];
        memory
            .read(&store, result_pointer as usize, &mut result_bytes)
            .map_err(|error| WasmTaskError::TaskAbi(error.to_string()))?;
        let result: WasmTaskResult = serde_json::from_slice(&result_bytes)
            .map_err(|error| WasmTaskError::TaskAbi(error.to_string()))?;
        result
            .validate_for(&invocation.task_instance)
            .map_err(WasmTaskError::TaskAbi)?;
        Ok(result)
    }

    pub fn run_task_export_verified_with_task_host(
        &self,
        wasm_or_wat: impl AsRef<[u8]>,
        expected_bundle_digest: &Digest,
        export: &str,
        invocation: &WasmTaskInvocation,
        host: Box<dyn WasmTaskHost>,
    ) -> Result<WasmTaskResult, WasmTaskError> {
        invocation.validate().map_err(WasmTaskError::TaskAbi)?;
        let module_bytes = Self::verified_module_bytes(wasm_or_wat, expected_bundle_digest)?;
        let encoded = serde_json::to_vec(invocation)
            .map_err(|error| WasmTaskError::TaskAbi(error.to_string()))?;
        let module = Module::new(&self.engine, &module_bytes).map_err(wasmtime_error)?;
        let linker = task_host_linker(&self.engine)?;
        let abort_signal = host.abort_signal();
        let debug_control = host.debug_control();
        let mut store = Store::new(
            &self.engine,
            WasmtimeTaskHostState {
                host,
                fatal_host_error: None,
                limits: task_store_limits(&self.runtime_limits),
                fuel_budget: FuelTokenBucket::new(&self.runtime_limits),
            },
        );
        store.limiter(|state| &mut state.limits);
        store
            .set_fuel(
                self.runtime_limits
                    .fuel_capacity()
                    .expect("validated fuel capacity"),
            )
            .map_err(wasmtime_error)?;
        defer_epoch_interruption(&mut store);
        let instance = linker
            .instantiate(&mut store, &module)
            .map_err(wasmtime_error)?;
        let memory = instance
            .get_memory(&mut store, "memory")
            .ok_or_else(|| WasmTaskError::TaskAbi("guest module exports no memory".to_owned()))?;
        let allocate = instance
            .get_typed_func::<u32, u32>(&mut store, "clusterflux_alloc_v1")
            .map_err(wasmtime_error)?;
        let input_length = u32::try_from(encoded.len())
            .map_err(|_| WasmTaskError::TaskAbi("task invocation is too large".to_owned()))?;
        let input_pointer = allocate
            .call(&mut store, input_length)
            .map_err(wasmtime_error)?;
        if input_pointer == 0 && input_length != 0 {
            return Err(WasmTaskError::TaskAbi(
                "guest refused task invocation allocation".to_owned(),
            ));
        }
        memory
            .write(&mut store, input_pointer as usize, &encoded)
            .map_err(|error| WasmTaskError::TaskAbi(error.to_string()))?;
        let task = instance
            .get_typed_func::<(u32, u32), u64>(&mut store, export)
            .map_err(wasmtime_error)?;
        let _abort_guard = abort_signal.as_ref().map(|signal| {
            EpochControlGuard::arm(&self.engine, &mut store, Arc::clone(signal), debug_control)
        });
        debug_control_trace("calling task ABI export");
        let packed = match task.call(&mut store, (input_pointer, input_length)) {
            Ok(packed) => packed,
            Err(error) => return Err(abort_error(&store, abort_signal.as_ref(), error)),
        };
        decode_guest_task_result(&store, &memory, packed, &invocation.task_instance)
    }

    fn run_i32_export_bytes(
        &self,
        wasm_or_wat: &[u8],
        export: &str,
        arg: i32,
    ) -> Result<i32, WasmTaskError> {
        let module = Module::new(&self.engine, wasm_or_wat).map_err(wasmtime_error)?;
        let linker = task_host_stub_linker(&self.engine)?;
        let mut store = Store::new(
            &self.engine,
            BasicStoreState {
                limits: task_store_limits(&self.runtime_limits),
            },
        );
        store.limiter(|state| &mut state.limits);
        store
            .set_fuel(
                self.runtime_limits
                    .fuel_capacity()
                    .expect("validated fuel capacity"),
            )
            .map_err(wasmtime_error)?;
        defer_epoch_interruption(&mut store);
        let instance = linker
            .instantiate(&mut store, &module)
            .map_err(wasmtime_error)?;
        let func = instance
            .get_typed_func::<i32, i32>(&mut store, export)
            .map_err(wasmtime_error)?;
        func.call(&mut store, arg).map_err(wasmtime_error)
    }

    pub fn freeze_resume_i32_export_probe(
        &self,
        wasm_or_wat: impl AsRef<[u8]>,
        export: &str,
        arg: i32,
    ) -> Result<WasmtimeDebugProbe, WasmTaskError> {
        let task = TaskInstanceId::from(export);
        let snapshot = Self::debug_i32_export_snapshot(wasm_or_wat.as_ref(), export, arg)?;
        let mut epoch = DebugEpoch::pause(
            ProcessId::from("wasmtime-debug-probe"),
            1,
            vec![DebugParticipant {
                task: task.clone(),
                name: export.to_owned(),
                kind: DebugParticipantKind::WasmTask,
                can_freeze: true,
                state: DebugRuntimeState::Running,
                stack_frames: snapshot.stack_frames.clone(),
                local_values: snapshot.local_values.clone(),
                task_args: vec![("arg".to_owned(), arg.to_string())],
                handles: Vec::new(),
                command_status: None,
                recent_output: Vec::new(),
            }],
        )
        .map_err(|err| WasmTaskError::Runtime(err.to_string()))?;
        let frozen_state = epoch
            .participant_state(&task)
            .cloned()
            .ok_or_else(|| WasmTaskError::Runtime("Wasm debug participant missing".to_owned()))?;
        let inspection = epoch
            .inspection(&task)
            .map_err(|err| WasmTaskError::Runtime(err.to_string()))?;
        epoch.continue_all();
        let resumed_state = epoch
            .participant_state(&task)
            .cloned()
            .ok_or_else(|| WasmTaskError::Runtime("Wasm debug participant missing".to_owned()))?;
        let result = self.run_i32_export(wasm_or_wat, export, arg)?;

        Ok(WasmtimeDebugProbe {
            task,
            frozen_state,
            resumed_state,
            result,
            stack_frames: inspection.stack_frames,
            local_values: inspection.local_values,
            wasm_function: snapshot.wasm_function,
            wasm_pc: snapshot.wasm_pc,
        })
    }

    fn debug_i32_export_snapshot(
        wasm_or_wat: &[u8],
        export: &str,
        arg: i32,
    ) -> Result<WasmtimeFrameSnapshot, WasmTaskError> {
        let engine = Self::debug_engine()?;
        let module = Module::new(&engine, wasm_or_wat).map_err(wasmtime_error)?;
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .map_err(|err| WasmTaskError::Runtime(format!("create debug runtime: {err:#}")))?;
        let state = runtime.block_on(async {
            let linker = task_host_stub_linker(&engine)?;
            let mut store = Store::new(&engine, WasmtimeDebugState::default());
            store.limiter(|state| &mut state.limits);
            store.set_debug_handler(WasmtimeLocalSnapshotHandler {
                export: export.to_owned(),
            });
            let instance = linker
                .instantiate_async(&mut store, &module)
                .await
                .map_err(wasmtime_error)?;
            if let Some(mut edit) = store.edit_breakpoints() {
                edit.single_step(true).map_err(wasmtime_error)?;
            }
            let func = instance
                .get_typed_func::<i32, i32>(&mut store, export)
                .map_err(wasmtime_error)?;
            let _ = func
                .call_async(&mut store, arg)
                .await
                .map_err(wasmtime_error)?;
            Ok::<_, WasmTaskError>(store.into_data())
        })?;

        state.snapshot.ok_or_else(|| {
            WasmTaskError::Runtime(
                "Wasmtime guest debug did not produce a frame-local snapshot".to_owned(),
            )
        })
    }

    fn verified_module_bytes(
        wasm_or_wat: impl AsRef<[u8]>,
        expected_bundle_digest: &Digest,
    ) -> Result<Vec<u8>, WasmTaskError> {
        let bytes = wasm_or_wat.as_ref();
        let actual = Digest::sha256(bytes);
        if &actual != expected_bundle_digest {
            return Err(WasmTaskError::BundleDigestMismatch {
                expected: expected_bundle_digest.clone(),
                actual,
            });
        }
        Ok(bytes.to_vec())
    }
}

fn wasmtime_error(err: wasmtime::Error) -> WasmTaskError {
    WasmTaskError::Runtime(format!("{err:?}"))
}

fn decode_guest_task_result(
    store: &Store<WasmtimeTaskHostState>,
    memory: &wasmtime::Memory,
    packed: u64,
    expected_task: &TaskInstanceId,
) -> Result<WasmTaskResult, WasmTaskError> {
    let result_pointer = packed as u32;
    let result_length = (packed >> 32) as u32;
    if result_length as usize > MAX_WASM_TASK_ENVELOPE_BYTES {
        return Err(WasmTaskError::TaskAbi(format!(
            "guest task result is {result_length} bytes; maximum is {MAX_WASM_TASK_ENVELOPE_BYTES}"
        )));
    }
    let mut result_bytes = vec![0_u8; result_length as usize];
    memory
        .read(store, result_pointer as usize, &mut result_bytes)
        .map_err(|error| WasmTaskError::TaskAbi(error.to_string()))?;
    let result: WasmTaskResult = serde_json::from_slice(&result_bytes)
        .map_err(|error| WasmTaskError::TaskAbi(error.to_string()))?;
    result
        .validate_for(expected_task)
        .map_err(WasmTaskError::TaskAbi)?;
    Ok(result)
}

#[cfg(test)]
mod tests;