oo-ide 0.0.4

//! Wasmtime instance lifecycle for a single extension.

use std::collections::HashMap;

use tokio::sync::mpsc::UnboundedSender;
use wasmtime::{AsContext, AsContextMut, Engine, Instance, Linker, Module, Store};

use oo_protocol::{HostRequest, Message, Request, Response, ResponsePayload};

use crate::operation::Operation;
use crate::prelude::*;

use super::op_bridge;

// ---------------------------------------------------------------------------
// Host state stored in the wasmtime Store
// ---------------------------------------------------------------------------

pub(crate) struct HostState {
    /// Channel back to the IDE event loop.
    pub op_tx: UnboundedSender<Vec<Operation>>,
    /// Extension-scoped config values.
    pub config: HashMap<String, String>,
    /// Project-detected languages (e.g. ["rust", "python"]).
    pub languages: Vec<String>,
}

// ---------------------------------------------------------------------------
// ExtensionInstance
// ---------------------------------------------------------------------------

pub(crate) struct ExtensionInstance {
    store: Store<HostState>,
    instance: Instance,
}

impl ExtensionInstance {
    /// Compile and instantiate a WASM extension module.
    pub(crate) fn new(
        engine: &Engine,
        wasm_bytes: &[u8],
        op_tx: UnboundedSender<Vec<Operation>>,
        config: HashMap<String, String>,
        languages: Vec<String>,
    ) -> Result<Self> {
        let module = Module::new(engine, wasm_bytes).map_err(|e| anyhow::anyhow!("compiling extension WASM: {}", e))?;

        let mut linker: Linker<HostState> = Linker::new(engine);

        // Single host-call function replacing the old __oo_operation /
        // __oo_log / __oo_get_config trio.
        linker
            .func_wrap(
                "oo",
                "__oo_host_call",
                |mut caller: wasmtime::Caller<'_, HostState>, ptr: i32, len: i32| {
                    let bytes = read_bytes(&mut caller, ptr, len).unwrap_or_default();
                    let msg: Message = match rmp_serde::from_slice(&bytes) {
                        Ok(m) => m,
                        Err(e) => {
                            log::warn!("extension sent invalid MessagePack: {}", e);
                            return;
                        }
                    };
                    if let Message::Request(req) = msg {
                        let id = req.id;
                        let result = handle_host_request(req, caller.data());
                        let response = Response { id, result };
                        send_response_to_wasm(&mut caller, response);
                    }
                },
            )
            .map_err(|e| anyhow::anyhow!("registering __oo_host_call: {}", e))?;

        // Legacy stub — kept so older extension binaries that still import
        // __oo_alloc via the "oo" module don't fail to link.
        linker
            .func_wrap(
                "oo",
                "__oo_alloc",
                |_caller: wasmtime::Caller<'_, HostState>, _size: i32| -> i32 { 0 },
            )
            .map_err(|e| anyhow::anyhow!("registering __oo_alloc stub: {}", e))?;

        let host_state = HostState { op_tx, config, languages };
        let mut store = Store::new(engine, host_state);

        let instance = linker
            .instantiate(&mut store, &module)
            .map_err(|e| anyhow::anyhow!("instantiating extension WASM: {}", e))?;

        // Call oo_init if the extension exports it.
        if let Ok(init_fn) = instance.get_typed_func::<(), ()>(&mut store, "oo_init") {
            init_fn.call(&mut store, ()).map_err(|e| anyhow::anyhow!("calling oo_init: {}", e))?;
        }

        Ok(Self { store, instance })
    }

    /// Call `detect_project()` — returns `true` if the extension is relevant.
    pub(crate) fn detect_project(&mut self) -> bool {
        let Ok(f) = self
            .instance
            .get_typed_func::<(), i32>(&mut self.store, "detect_project")
        else {
            return false;
        };
        matches!(f.call(&mut self.store, ()), Ok(1))
    }

    /// Call a named WASM export function with a YAML payload string.
    ///
    /// Used by the host to dispatch events and commands into the extension.
    pub(crate) fn call_export(&mut self, export_name: &str, payload_yaml: &str) -> Result<()> {
        let payload_bytes = payload_yaml.as_bytes();

        // Allocate memory in WASM for the payload.
        let alloc_fn = self
            .instance
            .get_typed_func::<i32, i32>(&mut self.store, "oo_alloc")
            .map_err(|e| anyhow!("extension must export oo_alloc: {}", e))?;
        let ptr = alloc_fn
            .call(&mut self.store, payload_bytes.len() as i32)
            .map_err(|e| anyhow::anyhow!("oo_alloc failed: {}", e))?;

        // Write payload into WASM memory.
        let mem = self
            .instance
            .get_memory(&mut self.store, "memory")
            .ok_or_else(|| anyhow!("extension must export 'memory'"))?;
        mem.write(&mut self.store, ptr as usize, payload_bytes)
            .map_err(|e| anyhow::anyhow!("writing payload to WASM memory: {}", e))?;

        // Call the export.
        let target = self
            .instance
            .get_typed_func::<(i32, i32), ()>(&mut self.store, export_name)
            .map_err(|e| anyhow!("extension does not export {:?}: {}", export_name, e))?;
        target
            .call(&mut self.store, (ptr, payload_bytes.len() as i32))
            .map_err(|e| anyhow::anyhow!("calling {}: {}", export_name, e))?;

        // Free the allocation.
        let free_fn = self
            .instance
            .get_typed_func::<(i32, i32), ()>(&mut self.store, "oo_free")
            .map_err(|e| anyhow!("extension must export oo_free: {}", e))?;
        free_fn
            .call(&mut self.store, (ptr, payload_bytes.len() as i32))
            .map_err(|e| anyhow::anyhow!("oo_free failed: {}", e))?;

        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Host request handler
// ---------------------------------------------------------------------------

fn handle_host_request(req: Request, state: &HostState) -> ResponsePayload {
    match req.payload {
        HostRequest::Log { message } => {
            log::info!("[ext] {}", message);
            ResponsePayload::None
        }

        HostRequest::GetConfig { key } => match state.config.get(&key).cloned() {
            Some(val) => ResponsePayload::String(val),
            None => ResponsePayload::Error("not found".into()),
        },

        HostRequest::HasLanguage { lang } => {
            // Check the project's detected languages vector provided in the
            // host state when the extension was instantiated.
            let found = state.languages.iter().any(|l| l == &lang);
            ResponsePayload::String(found.to_string())
        }

        HostRequest::Operation(op) => match op_bridge::to_operation(op) {
            Ok(ide_op) => {
                let _ = state.op_tx.send(vec![ide_op]);
                ResponsePayload::None
            }
            Err(e) => {
                log::warn!("extension op error: {}", e);
                ResponsePayload::Error(e.to_string())
            }
        },

    }
}

// ---------------------------------------------------------------------------
// Send a response back into the WASM instance
// ---------------------------------------------------------------------------

fn send_response_to_wasm(caller: &mut wasmtime::Caller<'_, HostState>, response: Response) {
    let msg = Message::Response(response);
    let bytes = match rmp_serde::to_vec(&msg) {
        Ok(b) => b,
        Err(e) => {
            log::warn!("failed to serialize response: {}", e);
            return;
        }
    };

    // Allocate WASM memory for the response.
    let alloc_fn = match caller.get_export("oo_alloc") {
        Some(wasmtime::Extern::Func(f)) => f,
        _ => {
            log::warn!("extension missing oo_alloc export");
            return;
        }
    };
    let mut results = [wasmtime::Val::I32(0)];
    if alloc_fn
        .call(
            caller.as_context_mut(),
            &[wasmtime::Val::I32(bytes.len() as i32)],
            &mut results,
        )
        .is_err()
    {
        return;
    }
    let ptr = match &results[0] {
        wasmtime::Val::I32(p) => *p as usize,
        _ => return,
    };

    // Write response bytes into WASM memory.
    let mem = match caller.get_export("memory") {
        Some(wasmtime::Extern::Memory(m)) => m,
        _ => return,
    };
    if mem.write(caller.as_context_mut(), ptr, &bytes).is_err() {
        return;
    }

    // Call oo_on_message(ptr, len) on the WASM instance.
    let on_msg = match caller.get_export("oo_on_message") {
        Some(wasmtime::Extern::Func(f)) => f,
        _ => return, // extension doesn't implement async callbacks — fine for fire-and-forget
    };
    let mut results = [];
    let _ = on_msg.call(
        caller.as_context_mut(),
        &[
            wasmtime::Val::I32(ptr as i32),
            wasmtime::Val::I32(bytes.len() as i32),
        ],
        &mut results,
    );

    // Free the allocation.
    let free_fn = match caller.get_export("oo_free") {
        Some(wasmtime::Extern::Func(f)) => f,
        _ => return,
    };
    let mut results = [];
    let _ = free_fn.call(
        caller.as_context_mut(),
        &[
            wasmtime::Val::I32(ptr as i32),
            wasmtime::Val::I32(bytes.len() as i32),
        ],
        &mut results,
    );
}

// ---------------------------------------------------------------------------
// Memory helpers
// ---------------------------------------------------------------------------

fn read_bytes(caller: &mut wasmtime::Caller<'_, HostState>, ptr: i32, len: i32) -> Result<Vec<u8>> {
    let mem = caller
        .get_export("memory")
        .and_then(|e| {
            if let wasmtime::Extern::Memory(m) = e {
                Some(m)
            } else {
                None
            }
        })
        .ok_or_else(|| anyhow!("extension must export 'memory'"))?;

    let data = mem.data(caller.as_context());
    let start = ptr as usize;
    let end = start + len as usize;
    data.get(start..end)
        .ok_or_else(|| anyhow!("out-of-bounds memory read"))
        .map(|s| s.to_vec())
}