binmod-mdk-core 0.1.4

use serde::{Serialize, Deserialize, de::DeserializeOwned};
use serde_json::{Value, Error as SerdeError};
use std::collections::HashMap;
use std::alloc::{alloc as std_alloc, dealloc as std_dealloc, Layout};


pub type FnResult<T> = Result<T, ModuleFnErr>;

#[unsafe(no_mangle)]
pub unsafe extern "C" fn guest_alloc(len: usize) -> *mut u8 {
    let align = std::mem::align_of::<usize>();
    let layout = Layout::from_size_align(len, align).unwrap();
    unsafe { std_alloc(layout) }
}

#[unsafe(no_mangle)]
pub unsafe extern "C" fn guest_dealloc(ptr: *mut u8, len: usize) {
    let align = std::mem::align_of::<usize>();
    let layout = Layout::from_size_align(len, align).unwrap();
    unsafe { std_dealloc(ptr, layout); }
}

pub fn pack_ptr(ptr: *const u8, len: usize) -> u64 {
    let ptr = ptr as u32;
    let len = len as u32;
    ((ptr as u64) << 32) | (len as u64)
}

pub fn unpack_ptr(packed: u64) -> (*const u8, usize) {
    let ptr = (packed >> 32) as *const u8;
    let len = (packed & 0xFFFFFFFF) as usize;
    (ptr, len)
}

pub fn serialize_to_ptr<T: Serialize>(result: T) -> Result<u64, SerdeError> {
    let bytes = serde_json::to_vec(&result)?;
    let ptr = pack_ptr(bytes.as_ptr(), bytes.len());
    // Ensure the bytes are not dropped
    std::mem::forget(bytes);

    Ok(ptr)
}

pub fn deserialize_from_ptr<T: DeserializeOwned>(input_ptr: u32, input_len: u32) -> Result<T, SerdeError> {
    let input_bytes = unsafe { std::slice::from_raw_parts(input_ptr as *const u8, input_len as usize) };
    serde_json::from_slice(input_bytes)
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ModuleFnInput {
    pub args: Option<Vec<Value>>,
    pub kwargs: Option<HashMap<String, Value>>,
}

impl ModuleFnInput {
    pub fn new() -> Self {
        Self {
            args: None,
            kwargs: None,
        }
    }

    pub fn get_arg<T: DeserializeOwned>(&self, index: usize, name: &str) -> Result<T, String> {
        if let Some(kwargs) = &self.kwargs {
            if let Some(value) = kwargs.get(name) {
                return serde_json::from_value(value.clone())
                    .map_err(|e| format!("Failed to parse argument {}: {}", name, e));
            }
        }

        if let Some(args) = &self.args {
            if index < args.len() {
                return serde_json::from_value(args[index].clone())
                    .map_err(|e| format!("Failed to parse argument {}: {}", index, e));
            }
        }

        Err(format!("Missing argument: {} in position {}", name, index))
    }

    pub fn get_args<T: DeserializeOwned>(&self) -> Result<Vec<T>, String> {
        if let Some(args) = &self.args {
            return serde_json::from_value(Value::Array(args.clone()))
                .map_err(|e| format!("Failed to parse arguments: {}", e));
        }
        Err("No arguments provided".into())
    }

    pub fn get_kwargs<T: DeserializeOwned>(&self) -> Result<HashMap<String, T>, String> {
        if let Some(kwargs) = &self.kwargs {
            return serde_json::from_value(Value::Object(
                kwargs
                    .iter()
                    .map(|(k, v)| (k.clone(), v.clone()))
                    .collect(),
            ))
                .map_err(|e| format!("Failed to parse keyword arguments: {}", e));
        }
        Err("No keyword arguments provided".into())
    }

    pub fn add_arg<T: Serialize>(&mut self, arg: T) -> Result<(), String> {
        if let Some(args) = &mut self.args {
            args.push(
                serde_json::to_value(arg)
                    .map_err(|e| format!("Failed to serialize argument: {}", e))?,
            );
        } else {
            self.args = Some(vec![
                serde_json::to_value(arg)
                    .map_err(|e| format!("Failed to serialize argument: {}", e))?,
            ]);
        }

        Ok(())
    }

    pub fn add_kwarg<T: Serialize>(&mut self, name: String, arg: T) -> Result<(), String> {
        if let Some(kwargs) = &mut self.kwargs {
            kwargs.insert(
                name, 
                serde_json::to_value(arg)
                    .map_err(|e| format!("Failed to serialize keyword argument: {}", e))?,
            );
        } else {
            let mut kwargs = HashMap::new();

            kwargs.insert(
                name, 
                serde_json::to_value(arg)
                    .map_err(|e| format!("Failed to serialize keyword argument: {}", e))?,
            );
            self.kwargs = Some(kwargs);
        }

        Ok(())
    }
}

impl Default for ModuleFnInput {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ModuleFnErr {
    #[serde(rename = "type")]
    pub error_type: String,
    pub message: String,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ModuleFnReturn<T: Serialize> {
    pub value: Option<T>,
}

impl<T: Serialize> ModuleFnReturn<T> {
    pub fn new(value: T) -> Self {
        Self { value: Some(value) }
    }

    pub fn empty() -> Self {
        Self { value: None }
    }
}

impl ModuleFnReturn<serde_json::Value> {
    pub fn new_serialized<T: Serialize>(value: T) -> Result<Self, SerdeError> {
        Ok(ModuleFnReturn::new(serde_json::to_value(value)?))
    }
}


#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(tag = "object")]
pub enum ModuleFnResult<T: Serialize> {
    #[serde(rename = "data")]
    Data(ModuleFnReturn<T>),
    #[serde(rename = "error")]
    Error(ModuleFnErr),
}