#![no_std]

#[unsafe(no_mangle)]
pub extern "C" fn add(left: u64, right: u64) -> u64 {
    left + right
}

#[unsafe(no_mangle)]
pub extern "C" fn subtract(left: u32, right: u32) -> u32 {
    left - right
}

>>> wasmtime --invoke add ./wasm_demo.wasm 1 2
3

(module $wasm_demo.wasm
  (type (;0;) (func (param i64 i64) (result i64)))
  (type (;1;) (func (param i32 i32) (result i32)))
  (func $add (type 0) (param i64 i64) (result i64)
    local.get 1
    local.get 0
    i64.add)
  (func $subtract (type 1) (param i32 i32) (result i32)
    local.get 0
    local.get 1
    i32.sub)
  (memory (;0;) 16)
  (global $__stack_pointer (mut i32) (i32.const 1048576))
  (global (;1;) i32 (i32.const 1048576))
  (global (;2;) i32 (i32.const 1048576))
  (export "memory" (memory 0))
  (export "add" (func $add))
  (export "subtract" (func $subtract))
  (export "__data_end" (global 1))
  (export "__heap_base" (global 2)))

(module $wasm_demo.wasm
  ...
)

(type (;0;) (func (param i64 i64) (result i64)))

(func $add (type 0) (param i64 i64) (result i64)
  local.get 1
  local.get 0
  i64.add)

(memory (;0;) 16)

(global $__stack_pointer (mut i32) (i32.const 1048576))

(export "memory" (memory 0))
(export "add" (func $add))
(export "subtract" (func $subtract))
(export "__data_end" (global 1))
(export "__heap_base" (global 2))

#[no_mangle]
pub extern "C" fn allocate(size: usize) -> *mut u8 {
    let mut buffer = Vec::with_capacity(size);
    let ptr = buffer.as_mut_ptr();
    core::mem::forget(buffer); // 阻止 Vec 在这里被 Drop
    ptr
}

#[no_mangle]
pub extern "C" fn deallocate(ptr: *mut u8, capacity: usize) {
    unsafe {
        let _ = Vec::from_raw_parts(ptr, 0, capacity);
    }
}

#[no_mangle]
pub extern "C" fn json_add_field(
    json_ptr: *const u8,
    json_len: usize,
    field_ptr: *const u8,
    field_len: usize,
) -> u64 {
  // 加载 host 传入的字符串，并解析
    let json_bytes = unsafe { slice::from_raw_parts(json_ptr, json_len) };
    let field_bytes = unsafe { slice::from_raw_parts(field_ptr, field_len) };
    let mut json_value: serde_json::Value = match serde_json::from_slice(json_bytes) {
        Ok(v) => v,
        Err(_) => return 0, // 错误处理
    };
    let field_value: serde_json::Value = match serde_json::from_slice(field_bytes) {
        Ok(v) => v,
        Err(_) => return 0, // 错误处理
    };

    // 执行添加字段的逻辑
    if let (Some(obj), Some(field_obj)) = (json_value.as_object_mut(), field_value.as_object()) {
        for (k, v) in field_obj {
            obj.insert(k.clone(), v.clone());
        }
    } else {
        return 0;
    }

    // 将修改后的 JSON 序列化回字符串
    let result_json = match serde_json::to_string(&json_value) {
        Ok(s) => s,
        Err(_) => return 0, // 错误处理
    };

    // 将结果字符串转换为字节向量并获取指针和长度
    let mut result_bytes = result_json.into_bytes();
    let len = result_bytes.len();
    let ptr = result_bytes.as_mut_ptr() as u64;

    // 泄漏内存，以便 Host 可以访问它
    core::mem::forget(result_bytes);

    // 将长度和指针打包成一个 u64 返回
    // 长度放在高 32 位，指针放在低 32 位
    // Q: WASM 支持 multi-value 返回，RUST 如何使用这个特性呢？
    (len as u64) << 32 | ptr
}

$ python run.py 
Initial JSON: {"name": "Alice", "age": 30}
Adding Fields: {"city": "New York", "occupation": "Engineer"}
Modified JSON: {"age":30,"city":"New York","name":"Alice","occupation":"Engineer"}

store = wasmtime.Store()
module = wasmtime.Module.from_file(store.engine, "wasms/rust.wasm")
instance = wasmtime.Instance(store, module, [])

# 获取 WASM 中的导出
exports = instance.exports(store)
wasm_memory = exports["memory"]
wasm_allocate = exports["allocate"]
wasm_deallocate = exports["deallocate"]
wasm_json_add_field = exports["json_add_field"]

def rust_string_to_python(ptr, length):
    data_bytes = wasm_memory.read(store, ptr, ptr + length)
    return data_bytes.decode('utf-8')

def python_string_to_wasm(s: str) -> tuple[int, int]:
    s_bytes = s.encode('utf-8')
    s_len = len(s_bytes)
    s_ptr = wasm_allocate(store, s_len)
    wasm_memory.write(store, s_bytes, s_ptr)
    return s_ptr, s_len

def process_json_with_wasm(original_json: str, field_to_add: str) -> str:
    json_ptr, json_len = python_string_to_wasm(original_json)
    field_ptr, field_len = python_string_to_wasm(field_to_add)
    # 调用函数
    result_u64 = wasm_json_add_field(store, json_ptr, json_len, field_ptr, field_len)
    # 释放 WASM 中的内存
    wasm_deallocate(store, json_ptr, json_len)
    wasm_deallocate(store, field_ptr, field_len)

    # 解析 WASM 函数返回的指针和长度
    if result_u64 == 0:
        raise ValueError("Wasm function returned an error (0)")

    result_len = result_u64 >> 32
    result_ptr = result_u64 & 0xFFFFFFFF
    result_json_str = rust_string_to_python(result_ptr, result_len)

    # 在 Rust 这部分内存是没有被释放的，需要手动释放
    wasm_deallocate(store, result_ptr, result_len)

    return result_json_str

if __name__ == "__main__":
    initial_json = '{"name": "Alice", "age": 30}'
    new_field = '{"city": "New York", "occupation": "Engineer"}'
    modified_json = process_json_with_wasm(initial_json, new_field)
    print(f"Modified JSON: {modified_json}")

async function runWasm() {
    // 1. 读取 Wasm 模块的二进制数据
    const wasmPath = path.join(__dirname, 'wasms', 'rust.wasm');
    const wasmBytes = fs.readFileSync(wasmPath);

    // 2. 创建 WebAssembly 内存实例
    // initial: 初始内存页数 (1页 = 64KB)
    // maximum: 最大内存页数 (可选，但推荐设置，防止内存无限增长)
    const memory = new WebAssembly.Memory({ initial: 1, maximum: 16 }); // 1页 = 64KB

    // 3. 自定义定义导入对象 (imports object)
    const importObject = {
        env: {
            memory: memory,
        }
    };

    // 4. 实例化 Wasm 模块
    const wasm = await WebAssembly.instantiate(wasmBytes, importObject);
    const exports = wasm.instance.exports;

    // 5. 获取导出的函数和内存
    const wasmAllocate = exports.allocate;
    const wasmDeallocate = exports.deallocate;
    const wasmJsonAddField = exports.json_add_field;
    const wasmMemory = exports.memory; // 直接获取导出的内存对象

    function writeStringToWasm(jsString) {
        const encoder = new TextEncoder('utf-8');
        const encodedBytes = encoder.encode(jsString);
        const byteLength = encodedBytes.length;
        const ptr = wasmAllocate(byteLength);
        const wasmByteView = new Uint8Array(wasmMemory.buffer);
        wasmByteView.set(encodedBytes, ptr);
        return { ptr, len: byteLength };
    }
    

    function readStringFromWasm(ptr, len) {
        const wasmByteView = new Uint8Array(wasmMemory.buffer, ptr, len);
        const decoder = new TextDecoder('utf-8');
        return decoder.decode(wasmByteView);
    }

    const initialJson = '{"name": "Alice", "age": 30}';
    const newField = '{"city": "New York", "occupation": "Engineer"}';

    console.log(`Initial JSON: ${initialJson}`);
    console.log(`Field to add: ${newField}`);

    try {
        // 写入原始 JSON 字符串到 Wasm 内存
        const { ptr: jsonPtr, len: jsonLen } = writeStringToWasm(initialJson);
        // 写入要添加的字段 JSON 字符串到 Wasm 内存
        const { ptr: fieldPtr, len: fieldLen } = writeStringToWasm(newField);

        // 调用 Wasm 函数
        // Rust 函数返回一个 u64，高 32 位是长度，低 32 位是指针
        const resultU64 = wasmJsonAddField(jsonPtr, jsonLen, fieldPtr, fieldLen);

        // 释放输入字符串的 Wasm 内存
        wasmDeallocate(jsonPtr, jsonLen);
        wasmDeallocate(fieldPtr, fieldLen);

        if (resultU64 === 0) {
            throw new Error("Wasm function returned an error (0)");
        }

        // 从 u64 中解析指针和长度
        const resultLen = Number(resultU64 >> 32n); // 使用 BigInt 操作 u64
        const resultPtr = Number(resultU64 & 0xFFFFFFFFn); // 使用 BigInt 操作 u64

        // 从 Wasm 内存中读取结果字符串
        const modifiedJson = readStringFromWasm(resultPtr, resultLen);
        console.log(`Modified JSON: ${modifiedJson}`);
        // 释放 Wasm 模块分配的内存
        wasmDeallocate(resultPtr, resultLen);

        console.log("Got JSON:", expectedJsonDict);
    } catch (e) {
        console.error(`An error occurred: ${e.message}`);
    }
}

runWasm();