// Copyright (C) 2026 Ordinary Labs, LLC.
//
// SPDX-License-Identifier: AGPL-3.0-only

use crate::Template;

use ordinary_utils::wasm::{read_mem, write_mem};

use bytes::{Bytes, BytesMut};
use parking_lot::Mutex;
use std::error::Error;
use std::sync::Arc;

use ordinary_config::TemplateFfiSerialization;
use wasi_common::{
    WasiCtx,
    sync::{WasiCtxBuilder, add_to_linker},
};
use wasmtime::{Caller, Linker, Memory, MemoryType, Store};

pub fn call(template: &Template, args: &Bytes) -> Result<Bytes, Box<dyn Error>> {
    let mut linker = Linker::new(&template.engine);
    add_to_linker(&mut linker, |s| s)?;

    let wasi = WasiCtxBuilder::new().build();
    let mut store = Store::new(&template.engine, wasi);

    let memory_ty = MemoryType::new(1, None);
    Memory::new(&mut store, memory_ty)?;

    let args = match template.config.ffi.serialization {
        TemplateFfiSerialization::FlexBufferVector => args.clone(),
    };

    linker.func_wrap(
        "env",
        "host_get_input",
        move |mut caller: Caller<'_, WasiCtx>| -> i64 {
            if let Some(res) = write_mem(&mut caller, args.as_ref()) {
                return res;
            }

            0
        },
    )?;

    let out = Arc::new(Mutex::new(Bytes::new()));
    let out_clone = Arc::clone(&out);

    linker.func_wrap(
        "env",
        "host_set_output",
        move |mut caller: Caller<'_, WasiCtx>, ptr: i32, len: i32| -> i32 {
            // todo: evaluate whether `zeroed` or `resize` is faster
            let mut buf = BytesMut::zeroed(len as usize);

            if read_mem(&mut caller, ptr, &mut buf).is_ok() {
                let mut out = out_clone.lock();
                *out = buf.into();

                drop(out);

                return 1;
            }

            0
        },
    )?;

    let lock = template.module.read();

    if let Some(module) = &*lock {
        linker.module(&mut store, "", module)?;
    }

    if linker
        .get_default(&mut store, "")?
        .typed::<(), ()>(&store)?
        .call(&mut store, ())
        .is_err()
    {
        return Err("wasm call failed".into());
    }

    let out = out.lock();
    Ok(out.clone())
}
