//! Intrinsics that represent helpers that implement async calls

use crate::intrinsics::component::ComponentIntrinsic;
use crate::intrinsics::p3::async_task::AsyncTaskIntrinsic;
use crate::intrinsics::{Intrinsic, RenderIntrinsicsArgs};
use crate::source::Source;

/// This enum contains intrinsics that implement async calls
#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq)]
#[allow(clippy::enum_variant_names)]
pub enum HostIntrinsic {
    /// Intrinsic used by the host to prepare trampoline calls
    ///
    /// # Host Intrinsic implementation function
    ///
    /// The function that implements this intrinsic has the following definition:
    ///
    /// ```ts
    /// type i32 = number;
    /// function prepareCall(memoryIdx: i32): boolean;
    /// ```
    ///
    PrepareCall,

    /// Intrinsic used by the host to signal the start of an async-lowered call
    ///
    /// This intrinsic signals the start of an Async call emitted by modules generated
    /// by wasmtime's Fused Adapter Compiler of Trampolines (FACT)
    ///
    /// This call indicates that an async-lowered import function is being called
    /// (that has either been async lifted or not on the callee side)
    ///
    /// This intrinsic returns a combination of the initial call status and optionally
    /// the handle to a waitable that should be awaited until it's time to (if necessary),
    /// packed into the same u32.
    ///
    /// # Host Intrinsic implementation function
    ///
    /// The function that implements this intrinsic has the following definition:
    ///
    /// ```ts
    /// type i32 = number;
    /// type u32 = number;
    /// type u64 = number;
    /// type Args = {
    ///     postReturnIdx: number | null,
    ///     getPostReturnFn: () => function | null,
    ///     callbackIdx: number | null,
    ///     getCallbackFn: () => function | null,
    /// };
    /// function asyncStartCall(args: Args, callee: function, paramCount: u32, resultCount: u32, flags: u32): u32;
    /// ```
    ///
    /// NOTE: args are gathered during Trampoline, and the rest of the arguments are fed in.
    ///
    AsyncStartCall,

    /// Start of an sync call emitted by modules generated by wasmtime's
    /// Fused Adapter Compiler of Trampolines (FACT)
    ///
    /// This call maps indicates a trampoline for a sync-lowered import
    /// of an async-lifted export, meaning that the *calling* component
    /// has sync lowered, but the callee has async lifted
    ///
    /// this intrinsic signals the start of an sync-lowered call
    /// of an async-lifted export.
    ///
    /// # Host Intrinsic implementation function
    ///
    /// The function that implements this intrinsic has the following definition:
    ///
    /// ```ts
    /// type i32 = number;
    /// function syncStartCall(callbackIdx: i32): boolean;
    /// ```
    ///
    SyncStartCall,

    /// Used for writing out an event's contents into component memory (`unpack_event`)
    StoreEventInComponentMemory,
}

impl HostIntrinsic {
    /// Retrieve dependencies for this intrinsic
    pub fn deps() -> &'static [&'static Intrinsic] {
        &[]
    }

    /// Retrieve global names for this intrinsic
    pub fn get_global_names() -> impl IntoIterator<Item = &'static str> {
        ["syncStartCall", "asyncStartCall", "prepareCall"]
    }

    /// Get the name for the intrinsic
    pub fn name(&self) -> &'static str {
        match self {
            Self::PrepareCall => "_prepareCall",
            Self::AsyncStartCall => "_asyncStartCall",
            Self::SyncStartCall => "_syncStartCall",
            Self::StoreEventInComponentMemory => "_storeEventInComponentMemory",
        }
    }

    /// Render an intrinsic to a string
    pub fn render(&self, output: &mut Source, _render_args: &RenderIntrinsicsArgs<'_>) {
        match self {
            // PrepareCall is called before an async-lowered import (from the host or another component)
            // is called from inside a component.
            //
            // It's primary function is to set up a Subtask which will be used by the callee.
            // see: https://github.com/WebAssembly/component-model/blob/main/design/mvp/Async.md#structured-concurrency
            //
            // NOTE: it's possible for PrepareCall to be combined with AsyncStartCall
            // in a future component model release.
            //
            // For toolchains that implement `PrepareCall` and `AsyncStartCall`/`AsyncCall`,
            // `startFn` and `returnFn` are functions that *perform* the relevant lifting and lowering
            // when a call is entered or exited.
            //
            // For example, Fused components produced by `wasm-tools compose` will have generated start and
            // return functions that perform the inter-component lifting and lowering that needs to happen.
            //
            // While lifting and lowering is still performed by this crate, we often do so to make handling
            // consistent and make results available to functions like `AsyncTask#onResolve` for any other machinery
            // that expects to have that information.
            //
            Self::PrepareCall => {
                let debug_log_fn = Intrinsic::DebugLog.name();
                let prepare_call_fn = Self::PrepareCall.name();
                let current_task_get_fn =
                    Intrinsic::AsyncTask(AsyncTaskIntrinsic::GetCurrentTask).name();
                let create_new_current_task_fn =
                    Intrinsic::AsyncTask(AsyncTaskIntrinsic::CreateNewCurrentTask).name();
                let set_global_current_task_meta_fn = Intrinsic::SetGlobalCurrentTaskMetaFn.name();

                output.push_str(&format!(
                  r#"
                    function {prepare_call_fn}(
                        memoryIdx,
                        getMemoryFn,
                        startFn,
                        returnFn,
                        callerComponentIdx,
                        calleeComponentIdx,
                        taskReturnTypeIdx,
                        calleeIsAsyncInt,
                        stringEncoding,
                        resultCountOrAsync,
                    ) {{
                        {debug_log_fn}('[{prepare_call_fn}()]', {{
                            memoryIdx,
                            callerComponentIdx,
                            calleeComponentIdx,
                            taskReturnTypeIdx,
                            calleeIsAsyncInt,
                            stringEncoding,
                            resultCountOrAsync,
                        }});
                        const argArray = [...arguments];

                        // value passed in *may* be as large as u32::MAX which may be mangled into -2
                        resultCountOrAsync >>>= 0;

                        let isAsync = false;
                        let hasResultPointer = false;
                        if (resultCountOrAsync === 2**32 - 1) {{
                            // prepare async with no result (u32::MAX)
                            isAsync = true;
                            hasResultPointer = false;
                        }} else if (resultCountOrAsync === 2**32 - 2) {{
                            // prepare async with result (u32::MAX - 1)
                            isAsync = true;
                            hasResultPointer = true;
                        }}

                        const currentCallerTaskMeta = {current_task_get_fn}(callerComponentIdx);
                        if (!currentCallerTaskMeta) {{
                            throw new Error('invalid/missing current task for caller during prepare call');
                        }}

                        const currentCallerTask = currentCallerTaskMeta.task;
                        if (!currentCallerTask) {{
                            throw new Error('unexpectedly missing task in meta for caller during prepare call');
                        }}

                        if (currentCallerTask.componentIdx() !== callerComponentIdx) {{
                            throw new Error(`task component idx [${{ currentCallerTask.componentIdx() }}] !== [${{ callerComponentIdx }}] (callee ${{ calleeComponentIdx }})`);
                        }}

                        let getCalleeParamsFn;
                        let resultPtr = null;
                        let directParamsArr;
                        if (hasResultPointer) {{
                            directParamsArr = argArray.slice(10, argArray.length - 1);
                            getCalleeParamsFn = () => directParamsArr;
                            resultPtr = argArray[argArray.length - 1];
                        }} else {{
                            directParamsArr = argArray.slice(10);
                            getCalleeParamsFn = () => directParamsArr;
                        }}

                        let encoding;
                        switch (stringEncoding) {{
                            case 0:
                                encoding = 'utf8';
                                break;
                            case 1:
                                encoding = 'utf16';
                                break;
                            case 2:
                                encoding = 'compact-utf16';
                                break;
                            default:
                                throw new Error(`unrecognized string encoding enum [${{stringEncoding}}]`);
                        }}

                        const subtask = currentCallerTask.createSubtask({{
                           componentIdx: callerComponentIdx,
                           parentTask: currentCallerTask,
                           isAsync,
                           callMetadata: {{
                              getMemoryFn,
                              memoryIdx,
                              resultPtr,
                              returnFn,
                              startFn,
                              stringEncoding,
                           }}
                        }});

                        const [newTask, newTaskID] = {create_new_current_task_fn}({{
                            componentIdx: calleeComponentIdx,
                            isAsync,
                            getCalleeParamsFn,
                            entryFnName: [
                                'task',
                                subtask.getParentTask().id(),
                                'subtask',
                                subtask.id(),
                                'new-prepared-async-task'
                            ].join('/'),
                            stringEncoding,
                        }});
                        newTask.setParentSubtask(subtask);
                        newTask.setReturnMemoryIdx(memoryIdx);
                        newTask.setReturnMemory(getMemoryFn);
                        subtask.setChildTask(newTask);

                        newTask.subtaskMeta = {{
                            subtask,
                            calleeComponentIdx,
                            callerComponentIdx,
                            getCalleeParamsFn,
                            stringEncoding,
                            isAsync,
                        }};

                        {set_global_current_task_meta_fn}({{
                            taskID: newTask.id(),
                            componentIdx: newTask.componentIdx(),
                        }});
                    }}
              "#
                ));
            }

            // AsyncStartCall is called just before an async-lowered import from component "A"
            // is called from inside component "B" (both host->guest and guest->guest calls).
            //
            // We don't need to do much here, because async `Task`s are created during execution of
            // CallWasm/CallInterface, rather than here.
            //
            Self::AsyncStartCall => {
                let debug_log_fn = Intrinsic::DebugLog.name();
                let async_start_call_fn = Self::AsyncStartCall.name();
                let get_or_create_async_state_fn =
                    Intrinsic::Component(ComponentIntrinsic::GetOrCreateAsyncState).name();
                let async_event_code_enum = Intrinsic::AsyncEventCodeEnum.name();
                let async_driver_loop_fn =
                    Intrinsic::AsyncTask(AsyncTaskIntrinsic::DriverLoop).name();
                let lookup_memories_for_component_fn = Intrinsic::LookupMemoriesForComponent.name();
                let current_component_idx_globals =
                    AsyncTaskIntrinsic::GlobalAsyncCurrentComponentIdxs.name();
                let get_current_task_fn =
                    Intrinsic::AsyncTask(AsyncTaskIntrinsic::GetCurrentTask).name();
                let with_global_current_task_meta_async_fn =
                    Intrinsic::WithGlobalCurrentTaskMetaFnAsync.name();
                let get_global_current_task_meta_fn = Intrinsic::GetGlobalCurrentTaskMetaFn.name();

                // TODO: lower here for non-zero param count
                // https://github.com/bytecodealliance/wasmtime/blob/69ef9afc11a2846248c9e94affca0223dbd033fc/crates/wasmtime/src/runtime/component/concurrent.rs#L1775
                //
                // NOTE: relying on the globals here for the currently executing  not ideal,
                // but given that prepare is synchronous and so is async start call, we should be OK
                // at least due to the execution order.
                //
                output.push_str(&format!(r#"
                    function {async_start_call_fn}(args, callee, paramCount, resultCount, flags) {{
                        const componentIdx = {current_component_idx_globals}.at(-1);

                        const globalTaskMeta = {get_global_current_task_meta_fn}(componentIdx);
                        if (!globalTaskMeta) {{ throw new Error('missing global current task globalTaskMeta'); }}
                        const taskID = globalTaskMeta.taskID;

                        {debug_log_fn}('[{async_start_call_fn}()] args', {{ args, componentIdx }});
                        const {{ getCallbackFn, callbackIdx, getPostReturnFn, postReturnIdx }} = args;

                        const preparedTaskMeta = {get_current_task_fn}(componentIdx, taskID);
                        if (!preparedTaskMeta) {{ throw new Error('unexpectedly missing current task'); }}

                        const preparedTask = preparedTaskMeta.task;
                        if (!preparedTask) {{ throw new Error('unexpectedly missing current task'); }}
                        if (!preparedTask.subtaskMeta) {{ throw new Error('missing subtask meta from prepare'); }}

                        const {{
                            subtask,
                            returnMemoryIdx,
                            getReturnMemoryFn,
                            callerComponentIdx,
                            calleeComponentIdx,
                            getCalleeParamsFn,
                            isAsync,
                            stringEncoding,
                        }} = preparedTask.subtaskMeta;
                        if (!subtask) {{ throw new Error("missing subtask from cstate during async start call"); }}
                        if (calleeComponentIdx !== preparedTask.componentIdx()) {{
                                throw new Error(`meta callee idx [${{calleeComponentIdx}}] != current task idx [${{preparedTask.componentIdx()}}] during async start call`);
                            }}
                        if (calleeComponentIdx !== componentIdx) {{
                            throw new Error("mismatched componentIdx for async start call (does not match prepare)");
                        }}

                        const argArray = [...arguments];

                        if (resultCount < 0 || resultCount > 1) {{ throw new Error('invalid/unsupported result count'); }}

                        const callbackFnName = 'callback_' + callbackIdx;
                        const callbackFn = getCallbackFn();
                        preparedTask.setCallbackFn(callbackFn, callbackFnName);
                        preparedTask.setPostReturnFn(getPostReturnFn());

                        if (resultCount < 0 || resultCount > 1) {{
                            throw new Error(`unsupported result count [${{ resultCount }}]`);
                        }}

                        const params = preparedTask.getCalleeParams();
                        if (paramCount !== params.length) {{
                            throw new Error(`unexpected callee param count [${{ params.length }}], {async_start_call_fn} invocation expected [${{ paramCount }}]`);
                        }}

                        const callerComponentState = {get_or_create_async_state_fn}(subtask.componentIdx());

                        const calleeComponentState = {get_or_create_async_state_fn}(preparedTask.componentIdx());
                        const calleeBackpressure = calleeComponentState.hasBackpressure();

                        // Set up a handler on subtask completion to lower results from the call into the caller's memory region.
                        //
                        // NOTE: during fused guest->guest calls this handler is triggered, but does not actually perform
                        // lowering manually, as fused modules provider helper functions that can
                        subtask.registerOnResolveHandler((res) => {{
                            {debug_log_fn}('[{async_start_call_fn}()] handling subtask result', {{ res, subtaskID: subtask.id() }});

                            let subtaskCallMeta = subtask.getCallMetadata();

                            // NOTE: in the case of guest -> guest async calls, there may be no memory/realloc present,
                            // as the host will intermediate the value storage/movement between calls.
                            //
                            // We can simply take the value and lower it as a parameter
                            if (subtaskCallMeta.memory || subtaskCallMeta.realloc) {{
                                throw new Error("call metadata unexpectedly contains memory/realloc for guest->guest call");
                            }}

                            const callerTask = subtask.getParentTask();
                            const calleeTask = preparedTask;
                            const callerMemoryIdx = callerTask.getReturnMemoryIdx();
                            const callerComponentIdx = callerTask.componentIdx();

                            // If a helper function was provided we are likely in a fused guest->guest call,
                            // and the result will be delivered (lift/lowered) via helper function
                            if (subtaskCallMeta && subtaskCallMeta.returnFn) {{
                                {debug_log_fn}('[{async_start_call_fn}()] return function present while handling subtask result, returning early (skipping lower)');

                                // TODO: centralize calling of returnFn to *one place* (if possible)
                                if (subtaskCallMeta.returnFnCalled) {{ return; }}

                                subtaskCallMeta.returnFn.apply(null, [subtaskCallMeta.resultPtr]);
                                return;
                            }}

                            // If there is no where to lower the results, exit early
                            if (!subtaskCallMeta.resultPtr) {{
                                {debug_log_fn}('[{async_start_call_fn}()] no result ptr during subtask result handling, returning early (skipping lower)');
                                return;
                            }}

                            let callerMemory;
                            if (callerMemoryIdx !== null && callerMemoryIdx !== undefined) {{
                                callerMemory = {lookup_memories_for_component_fn}({{ componentIdx: callerComponentIdx, memoryIdx: callerMemoryIdx }});
                            }} else {{
                                const callerMemories = {lookup_memories_for_component_fn}({{ componentIdx: callerComponentIdx }});
                                if (callerMemories.length !== 1) {{ throw new Error(`unsupported amount of caller memories`); }}
                                callerMemory = callerMemories[0];
                            }}

                            if (!callerMemory) {{
                                {debug_log_fn}('[{async_start_call_fn}()] missing memory', {{ subtaskID: subtask.id(), res }});
                                throw new Error(`missing memory for to guest->guest call result (subtask [${{subtask.id()}}])`);
                            }}

                            const lowerFns = calleeTask.getReturnLowerFns();
                            if (!lowerFns || lowerFns.length === 0) {{
                                {debug_log_fn}('[{async_start_call_fn}()] missing result lower metadata for guest->guest call', {{ subtaskID: subtask.id() }});
                                throw new Error(`missing result lower metadata for guest->guest call (subtask [${{subtask.id()}}])`);
                            }}

                            if (lowerFns.length !== 1) {{
                                {debug_log_fn}('[{async_start_call_fn}()] only single result reportetd for guest->guest call', {{ subtaskID: subtask.id() }});
                                throw new Error(`only single result supported for guest->guest calls (subtask [${{subtask.id()}}])`);
                            }}

                            {debug_log_fn}('[{async_start_call_fn}()] lowering results', {{ subtaskID: subtask.id() }});
                            lowerFns[0]({{
                                realloc: undefined,
                                memory: callerMemory,
                                vals: [res],
                                storagePtr: subtaskCallMeta.resultPtr,
                                componentIdx: callerComponentIdx,
                                stringEncoding: subtaskCallMeta.stringEncoding,
                            }});

                        }});

                        subtask.setOnProgressFn(() => {{
                            subtask.setPendingEvent(() => {{
                                if (subtask.isResolved()) {{ subtask.deliverResolve(); }}
                                const event = {{
                                    code: {async_event_code_enum}.SUBTASK,
                                    payload0: subtask.waitableRep(),
                                    payload1: subtask.getStateNumber(),
                                }};
                                return event;
                            }});
                        }});

                        // Start the (event) driver loop that will resolve the task
                        queueMicrotask(async () => {{
                            let startRes = subtask.onStart({{ startFnParams: params }});
                            startRes = Array.isArray(startRes) ? startRes : [startRes];

                            await calleeComponentState.suspendTask({{
                                task: preparedTask,
                                readyFn: () => !calleeComponentState.isExclusivelyLocked(),
                            }});

                            const started = await preparedTask.enter();
                            if (!started) {{
                                {debug_log_fn}('[{async_start_call_fn}()] task failed early', {{
                                    taskID: preparedTask.id(),
                                    subtaskID: subtask.id(),
                                }});
                                throw new Error("task failed to start");
                                return;
                            }}

                            let callbackResult;
                            try {{
                                let jspiCallee = WebAssembly.promising(callee);
                                callbackResult = await {with_global_current_task_meta_async_fn}({{
                                    taskID: preparedTask.id(),
                                    componentIdx: preparedTask.componentIdx(),
                                    fn: () => {{
                                        return jspiCallee.apply(null, startRes);
                                    }}
                                }});
                            }} catch(err) {{
                                {debug_log_fn}("[{async_start_call_fn}()] initial subtask callee run failed", err);
                                // NOTE: a good place to rejectt the parent task, if rejection API is enabled
                                // subtask.reject(err);
                                // subtask.getParentTask().reject(err);

                                subtask.getParentTask().setErrored(err);

                                return;
                            }}

                            // If there was no callback function, we're dealing with a sync function
                            // that was lifted as async without one, there is only the callee.
                            if (!callbackFn) {{
                                {debug_log_fn}("[{async_start_call_fn}()] no callback, resolving w/ callee result", {{
                                    taskID: preparedTask.id(),
                                    componentIdx: preparedTask.componentIdx(),
                                    preparedTask,
                                    stateNumber: preparedTask.taskState(),
                                    isResolved: preparedTask.isResolved(),
                                    callbackFn,
                                }});
                                preparedTask.resolve([callbackResult]);
                                return;
                            }}

                            let fnName = callbackFn.fnName;
                            if (!fnName) {{
                                fnName = [
                                    '<task ',
                                    subtask.parentTaskID(),
                                    '/subtask ',
                                    subtask.id(),
                                    '/task ',
                                    preparedTask.id(),
                                    '>',
                                ].join("");
                            }}

                            try {{
                                {debug_log_fn}("[{async_start_call_fn}()] starting driver loop", {{
                                    fnName,
                                    componentIdx: preparedTask.componentIdx(),
                                    subtaskID: subtask.id(),
                                    childTaskID: subtask.childTaskID(),
                                    parentTaskID: subtask.parentTaskID(),
                                }});

                                await {async_driver_loop_fn}({{
                                    componentState: calleeComponentState,
                                    task: preparedTask,
                                    fnName,
                                    isAsync: true,
                                    callbackResult,
                                    resolve,
                                    reject
                                }});
                            }} catch (err) {{
                                {debug_log_fn}("[AsyncStartCall] drive loop call failure", {{ err }});
                            }}

                        }});

                        const subtaskState = subtask.getStateNumber();
                        if (subtaskState < 0 || subtaskState > 2**5) {{
                            throw new Error('invalid subtask state, out of valid range');
                        }}

                        {debug_log_fn}('[{async_start_call_fn}()] returning subtask rep & state', {{
                            subtask: {{
                                rep: subtask.waitableRep(),
                                state: subtaskState,
                            }}
                        }});

                        return Number(subtask.waitableRep()) << 4 | subtaskState;
                    }}
                "#));
            }

            Self::SyncStartCall => {
                let debug_log_fn = Intrinsic::DebugLog.name();
                let sync_start_call_fn = Self::SyncStartCall.name();
                output.push_str(&format!(
                    "
                    function {sync_start_call_fn}(callbackIdx) {{
                        {debug_log_fn}('[{sync_start_call_fn}()] args', {{ callbackIdx }});
                        throw new Error('synchronous start call not implemented!');
                    }}
                "
                ));
            }

            Self::StoreEventInComponentMemory => {
                let debug_log_fn = Intrinsic::DebugLog.name();
                let store_event_in_component_memory_fn = Self::StoreEventInComponentMemory.name();
                output.push_str(&format!(
                    r#"
                    function {store_event_in_component_memory_fn}(args) {{
                        {debug_log_fn}('[{store_event_in_component_memory_fn}()] args', args);
                        const {{ memory, ptr, event }} = args;

                        if (!memory) {{ throw new Error('unexpectedly missing memory'); }}
                        if (ptr === undefined || ptr === null) {{ throw new Error('unexpectedly missing pointer'); }}
                        if (!event) {{ throw new Error('event object missing'); }}
                        if (event.code === undefined) {{ throw new Error('invalid event object, missing code'); }}
                        if (event.payload0 === undefined) {{ throw new Error('invalid event object, missing payload0'); }}
                        if (event.payload1 === undefined) {{ throw new Error('invalid event object, missing payload1'); }}

                        const dv = new DataView(memory.buffer);
                        dv.setUint32(ptr, event.payload0, true);
                        dv.setUint32(ptr + 4, event.payload1, true);

                        return event.code;
                    }}
                    "#
                ));
            }
        }
    }
}
