luminarys-sdk 0.1.0

//! WASM ABI exports, handler registration, and method descriptors.
//!
//! Use [`skill_main!`] to emit all required WASM exports.

use crate::types::{InvokeRequest, InvokeResponse, VERSION};
use std::cell::RefCell;

/// SDK version packed as u32.
pub const SDK_VERSION: u32 = (0 << 16) | (1 << 8) | 0; // 0.1.0

// ── Handler type ──────────────────────────────────────────────────────────────

/// The function signature every skill handler must implement.
pub type Handler = fn(InvokeRequest) -> InvokeResponse;

// ── Global registry (single-threaded WASM) ────────────────────────────────────

thread_local! {
    static HANDLER: RefCell<Option<Handler>> = RefCell::new(None);
    static METHODS: RefCell<Vec<MethodInfo>> = RefCell::new(Vec::new());
    static EXTRA: RefCell<Vec<(String, Handler)>> = RefCell::new(Vec::new());
    static REQUIREMENTS: RefCell<Vec<RequirementInfo>> = RefCell::new(Vec::new());
    static IDENTITY: RefCell<(String, String, String, String)> = RefCell::new(Default::default());
    /// GC-safe buffer for results returned to the host via `skill_handle`.
    /// Kept alive until the next `skill_handle` call.
    static RESULT_BUF: RefCell<Vec<u8>> = RefCell::new(Vec::new());
}

/// Set the skill's identity metadata from annotations.
pub fn set_skill_identity(id: &str, name: &str, version: &str, description: &str) {
    IDENTITY.with(|i| {
        *i.borrow_mut() = (id.into(), name.into(), version.into(), description.into());
    });
}

// ── Registration ──────────────────────────────────────────────────────────────

/// Register the main dispatch handler and method descriptors.
pub fn register(h: Option<Handler>, methods: Vec<MethodInfo>) {
    if let Some(handler) = h {
        HANDLER.with(|slot| *slot.borrow_mut() = Some(handler));
    }
    METHODS.with(|m| {
        let mut list = m.borrow_mut();
        for info in &methods {
            if let Some(fn_) = info.handler_fn {
                EXTRA.with(|e| e.borrow_mut().push((info.name.clone(), fn_)));
            }
        }
        list.extend(methods);
    });
}

/// Register additional methods without replacing the generated dispatcher.
pub fn register_methods(methods: Vec<MethodInfo>) {
    register(None, methods);
}

/// Register skill requirements (from @skill:require annotations).
pub fn register_requirements(reqs: Vec<RequirementInfo>) {
    REQUIREMENTS.with(|r| r.borrow_mut().extend(reqs));
}

// ── Dispatch ──────────────────────────────────────────────────────────────────

/// Dispatch an [`InvokeRequest`] to the registered handler.
pub fn call_handler(req: InvokeRequest) -> InvokeResponse {
    let found = EXTRA.with(|e| {
        e.borrow()
            .iter()
            .find(|(name, _)| name == &req.method)
            .map(|(_, h)| *h)
    });
    if let Some(h) = found {
        return h(req);
    }
    HANDLER.with(|slot| {
        if let Some(h) = *slot.borrow() {
            h(req)
        } else {
            error_response("no handler registered — call register() or use skill_main!()")
        }
    })
}

/// Return the [`SkillDescriptor`] for the currently registered methods.
pub fn get_descriptor() -> SkillDescriptor {
    let (id, name, ver, desc) = IDENTITY.with(|i| i.borrow().clone());
    SkillDescriptor {
        version: VERSION,
        skill_id: id,
        skill_name: name,
        skill_version: ver,
        description: desc,
        methods: METHODS.with(|m| m.borrow().clone()),
        requirements: REQUIREMENTS.with(|r| r.borrow().clone()),
        sdk_version: SDK_VERSION,
    }
}

// ── skill_alloc_impl ──────────────────────────────────────────────────────────

/// Allocate `size` bytes for the host to write the request into.
/// Returns a pointer to the data region.
/// The host MUST call `skill_free(ptr)` with the same pointer after
/// `skill_handle` has been called.
pub fn skill_alloc_impl(size: u32) -> u32 {
    if size == 0 {
        return 0;
    }
    let total = size as usize + 8;
    let mut buf: Vec<u8> = Vec::with_capacity(total);
    // SAFETY: we just allocated `total` bytes; initialising to 0.
    unsafe { buf.set_len(total) };
    // Write total allocation length into the 8-byte header.
    buf[..8].copy_from_slice(&(total as u64).to_le_bytes());
    let data_ptr = unsafe { buf.as_mut_ptr().add(8) };
    // Transfer ownership to the host; skill_free will reconstruct and drop.
    std::mem::forget(buf);
    data_ptr as u32
}

// ── skill_free_impl ───────────────────────────────────────────────────────────

/// Free the buffer previously returned by `skill_alloc`.
///
/// Reads the 8-byte length header at `ptr - 8`, reconstructs the `Vec<u8>`,
/// and drops it — returning memory to Rust's allocator with no leak.
///
/// # Safety
/// `ptr` MUST be a value previously returned by `skill_alloc_impl`.
/// Calling with any other pointer is undefined behaviour.
pub fn skill_free_impl(ptr: u32) {
    if ptr == 0 {
        return;
    }
    // SAFETY: ptr was produced by skill_alloc_impl which stored the total
    // length in the 8 bytes immediately before the data pointer.
    unsafe {
        let header_ptr = (ptr as usize - 8) as *mut u8;
        // Read the 8-byte little-endian total length.
        let total_len =
            u64::from_le_bytes(*(header_ptr as *const [u8; 8])) as usize;
        // Reconstruct the Vec from (header_ptr, total_len, total_len) and drop it.
        // This runs the Vec's drop glue which calls the global allocator's `dealloc`.
        drop(Vec::from_raw_parts(header_ptr, total_len, total_len));
    }
}

// ── skill_handle_impl ─────────────────────────────────────────────────────────

/// Deserialise the request, dispatch to the registered handler,
/// serialise the response and return the result.
pub fn skill_handle_impl(req_ptr: u32, req_len: u32) -> u64 {
    // SAFETY: host guarantees [req_ptr, req_ptr+req_len) is valid for the
    // duration of this call. We copy immediately via from_slice, so no
    // dangling reference can outlive this scope.
    let bytes = unsafe {
        std::slice::from_raw_parts(req_ptr as *const u8, req_len as usize)
    };
    let req: InvokeRequest = match rmp_serde::from_slice(bytes) {
        Ok(r) => r,
        Err(e) => {
            return write_result(marshal_error(&format!("unmarshal InvokeRequest: {e}")));
        }
    };
    let mut resp = call_handler(req);
    resp.version = VERSION;
    let encoded = match rmp_serde::to_vec_named(&resp) {
        Ok(b) => b,
        Err(e) => marshal_error(&format!("marshal InvokeResponse: {e}")),
    };
    write_result(encoded)
}

// ── skill_describe_impl ───────────────────────────────────────────────────────

/// Serialise the registered [`SkillDescriptor`] and return a packed pointer.
pub fn skill_describe_impl() -> u64 {
    let desc = get_descriptor();
    let encoded = rmp_serde::to_vec_named(&desc).unwrap_or_default();
    write_result(encoded)
}

// ── Internal helpers ──────────────────────────────────────────────────────────

/// Store `bytes` in the thread-local result buffer and return the packed ptr.
///
/// The buffer is kept alive in `RESULT_BUF` until the next call, which is
/// safe because the host reads the result synchronously before returning
/// to the WASM module.
fn write_result(bytes: Vec<u8>) -> u64 {
    RESULT_BUF.with(|buf| {
        *buf.borrow_mut() = bytes;
        let b = buf.borrow();
        if b.is_empty() {
            return 0;
        }
        let ptr = b.as_ptr() as u64;
        let len = b.len() as u64;
        (ptr << 32) | len
    })
}

// ── skill_main! macro ─────────────────────────────────────────────────────────

/// Emit the four required WASM ABI exports for this skill.
///
/// Place at the crate root (usually `src/lib.rs`):
///
/// ```rust,ignore
/// use luminarys_sdk::prelude::*;
/// skill_main!();
/// ```
///
/// Expands to four `#[no_mangle] extern "C"` functions:
/// `skill_alloc`, `skill_free`, `skill_handle`, `skill_describe`.
///
/// If you need custom initialisation (e.g. lazy `register()` calls), write
/// the exports manually and call `ensure_init()` before delegating to the
/// `_impl` functions — see `examples/echo-skill` for the pattern.
#[macro_export]
macro_rules! skill_main {
    () => {
        /// Allocate a buffer for the host to write the request.
        #[no_mangle]
        pub extern "C" fn skill_alloc(size: u32) -> u32 {
            $crate::entrypoint::skill_alloc_impl(size)
        }

        /// Free the buffer previously returned by `skill_alloc`.
        #[no_mangle]
        pub extern "C" fn skill_free(ptr: u32) {
            $crate::entrypoint::skill_free_impl(ptr);
        }

        /// Dispatch the request and return the packed result pointer.
        #[no_mangle]
        pub extern "C" fn skill_handle(req_ptr: u32, req_len: u32) -> u64 {
            $crate::entrypoint::skill_handle_impl(req_ptr, req_len)
        }

        /// Return the encoded SkillDescriptor.
        #[no_mangle]
        pub extern "C" fn skill_describe() -> u64 {
            $crate::entrypoint::skill_describe_impl()
        }
    };
}

// ── MethodInfo ────────────────────────────────────────────────────────────────

/// Describes a single exported method including its parameter schema.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Default)]
#[serde(default)]
pub struct MethodInfo {
    #[serde(rename = "name")]
    pub name: String,
    #[serde(rename = "description")]
    pub description: String,
    #[serde(rename = "params", default, skip_serializing_if = "Vec::is_empty")]
    pub params: Vec<ParamInfo>,
    #[serde(rename = "mcp_hidden", default, skip_serializing_if = "is_false")]
    pub mcp_hidden: bool,
    #[serde(rename = "private_callback", default, skip_serializing_if = "is_false")]
    pub private_callback: bool,
    #[serde(skip)]
    pub(crate) handler_fn: Option<Handler>,
}

/// Start building a [`MethodInfo`].
pub fn method(name: impl Into<String>, description: impl Into<String>) -> MethodInfo {
    MethodInfo {
        name: name.into(),
        description: description.into(),
        params: vec![],
        mcp_hidden: false,
        private_callback: false,
        handler_fn: None,
    }
}

impl MethodInfo {
    /// Add a parameter descriptor.
    pub fn param(mut self, name: impl Into<String>, p: ParamInfo) -> Self {
        let mut pi = p;
        pi.name = name.into();
        self.params.push(pi);
        self
    }
    /// Hide from MCP `tools/list`; still callable via `call_module`.
    pub fn set_internal(mut self) -> Self { self.mcp_hidden = true; self }
    /// Private callback — hidden from MCP AND only callable by this skill.
    pub fn callback(mut self) -> Self {
        self.mcp_hidden = true;
        self.private_callback = true;
        self
    }
    /// Attach an inline handler.
    pub fn handle(mut self, f: Handler) -> Self { self.handler_fn = Some(f); self }
}

// ── ParamInfo ─────────────────────────────────────────────────────────────────

/// Describes one parameter of a method.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Default)]
#[serde(default)]
pub struct ParamInfo {
    #[serde(rename = "name", default)]
    pub name: String,
    #[serde(rename = "description", default, skip_serializing_if = "String::is_empty")]
    pub description: String,
    #[serde(rename = "type")]
    pub type_val: String,
    #[serde(rename = "required", default, skip_serializing_if = "is_false")]
    pub is_required: bool,
    #[serde(rename = "enum", default, skip_serializing_if = "Vec::is_empty")]
    pub enum_vals: Vec<String>,
    #[serde(rename = "items", default, skip_serializing_if = "Option::is_none")]
    pub items: Option<Box<ParamInfo>>,
}

impl ParamInfo {
    pub fn required(mut self) -> Self { self.is_required = true; self }
    pub fn desc(mut self, d: impl Into<String>) -> Self { self.description = d.into(); self }
    pub fn enum_vals(mut self, vals: Vec<impl Into<String>>) -> Self {
        self.enum_vals = vals.into_iter().map(|v| v.into()).collect();
        self
    }
}

fn mk_param(type_val: &str) -> ParamInfo {
    ParamInfo {
        name: String::new(),
        description: String::new(),
        type_val: type_val.into(),
        is_required: false,
        enum_vals: vec![],
        items: None,
    }
}

pub fn type_string() -> ParamInfo { mk_param("string") }
pub fn type_int()    -> ParamInfo { mk_param("integer") }
pub fn type_number() -> ParamInfo { mk_param("number") }
pub fn type_bool()   -> ParamInfo { mk_param("boolean") }
pub fn type_object() -> ParamInfo { mk_param("object") }
pub fn type_array(items: ParamInfo) -> ParamInfo {
    ParamInfo { type_val: "array".into(), items: Some(Box::new(items)), ..mk_param("array") }
}

// ── SkillDescriptor ───────────────────────────────────────────────────────────

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Default)]
#[serde(default)]
pub struct SkillDescriptor {
    #[serde(rename = "version")]
    pub version: u32,
    #[serde(rename = "skill_id", default, skip_serializing_if = "String::is_empty")]
    pub skill_id: String,
    #[serde(rename = "skill_name", default, skip_serializing_if = "String::is_empty")]
    pub skill_name: String,
    #[serde(rename = "skill_version", default, skip_serializing_if = "String::is_empty")]
    pub skill_version: String,
    #[serde(rename = "description", default, skip_serializing_if = "String::is_empty")]
    pub description: String,
    #[serde(rename = "methods")]
    pub methods: Vec<MethodInfo>,
    #[serde(rename = "requirements", default, skip_serializing_if = "Vec::is_empty")]
    pub requirements: Vec<RequirementInfo>,
    /// SDK version packed as u32: (major << 16) | (minor << 8) | patch.
    #[serde(rename = "sdk_version", default)]
    pub sdk_version: u32,
}

/// Declares one permission the skill expects from its manifest.
/// Populated by codegen from @skill:require annotations.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Default)]
pub struct RequirementInfo {
    #[serde(rename = "kind")]
    pub kind: String,
    #[serde(rename = "pattern", default, skip_serializing_if = "String::is_empty")]
    pub pattern: String,
    #[serde(rename = "mode", default, skip_serializing_if = "String::is_empty")]
    pub mode: String,
}

// ── Convenience constructors ──────────────────────────────────────────────────

/// Build a successful [`InvokeResponse`] with an encoded payload.
pub fn ok_response<T: serde::Serialize>(value: &T) -> InvokeResponse {
    let payload = rmp_serde::to_vec_named(value).unwrap_or_default();
    InvokeResponse { version: VERSION, payload, ..Default::default() }
}

/// Build an error [`InvokeResponse`].
pub fn error_response(msg: &str) -> InvokeResponse {
    InvokeResponse { version: VERSION, error: msg.to_owned(), ..Default::default() }
}

/// Encode an error-only [`InvokeResponse`] to bytes.
pub fn marshal_error(msg: &str) -> Vec<u8> {
    rmp_serde::to_vec_named(&error_response(msg)).unwrap_or_default()
}

fn is_false(v: &bool) -> bool { !v }