src/lua/lua_promise.cpp

#ifdef WITH_RIVE_SCRIPTING
#include "rive/lua/rive_lua_libs.hpp"

#include "lua.h"
#include "lualib.h"

#include <cstring>

namespace rive
{

// ---------------------------------------------------------------------------
// ScriptedPromise – a lightweight Promise/A+ for Luau with cancellation
// ---------------------------------------------------------------------------
//
// States: Pending → Fulfilled | Rejected | Cancelled  (one-way, immutable)
//
// Cancellation follows the Roblox model:
//   - cancel() propagates DOWN to consumers (chained promises)
//   - cancel() propagates UP to parent when ALL consumers are cancelled
//   - An onCancel hook fires to clean up in-flight work (e.g. image decode)
//
// Parent/consumer refs form a tree. To break the cycle (parent ↔ consumer),
// all refs are released eagerly on any terminal transition via
// promise_cleanupLinks().

// ── helpers ────────────────────────────────────────────────────────────────

static void promise_notifyCallbacks(lua_State* L, ScriptedPromise* p);
static void handleCoroutineCompletion(lua_State* L,
                                      lua_State* co,
                                      int status,
                                      int coIdx,
                                      int asyncPromiseIdx);

// Release parent/consumer/onCancel refs after a terminal transition.
static void promise_cleanupLinks(lua_State* L, ScriptedPromise* p)
{
    if (p->m_parentRef != LUA_NOREF)
    {
        lua_unref(L, p->m_parentRef);
        p->m_parentRef = LUA_NOREF;
    }
    for (int ref : p->m_consumerRefs)
    {
        if (ref != LUA_NOREF)
            lua_unref(L, ref);
    }
    p->m_consumerRefs.clear();
    if (p->m_onCancelRef != LUA_NOREF)
    {
        lua_unref(L, p->m_onCancelRef);
        p->m_onCancelRef = LUA_NOREF;
    }
}

// ── resolve / reject / cancel ─────────────────────────────────────────────

void ScriptedPromise::resolve(lua_State* L, int valueIdx)
{
    if (m_promiseState != PromiseState::Pending)
        return;

    // Promise/A+ flattening: if the resolved value is itself a Promise,
    // adopt its state instead of resolving with the Promise object.
    auto* inner = lua_torive<ScriptedPromise>(L, valueIdx, true);
    if (inner && inner != this)
    {
        if (inner->isFulfilled())
        {
            lua_rawgeti(L, LUA_REGISTRYINDEX, inner->resultRef());
            resolve(L, lua_gettop(L));
            lua_pop(L, 1);
            return;
        }
        if (inner->isRejected())
        {
            lua_rawgeti(L, LUA_REGISTRYINDEX, inner->resultRef());
            reject(L, lua_gettop(L));
            lua_pop(L, 1);
            return;
        }
        if (inner->isCancelled())
        {
            cancel(L);
            return;
        }
        // Still pending — adopt by chaining. Handled by the overload
        // that has a selfRef. Direct resolve() callers that don't have
        // a ref to self can't adopt pending promises; assert that this
        // path is unreachable in practice.
        assert(false && "resolve() with pending inner promise requires "
                        "selfRef — use the overload");
        return;
    }

    m_promiseState = PromiseState::Fulfilled;
    lua_pushvalue(L, valueIdx);
    m_resultRef = lua_ref(L, -1);
    lua_pop(L, 1);
    promise_notifyCallbacks(L, this);
    promise_cleanupLinks(L, this);
}

void ScriptedPromise::resolve(lua_State* L, int valueIdx, int selfRef)
{
    if (m_promiseState != PromiseState::Pending)
        return;

    // Promise/A+ flattening with self registry ref for pending adoption.
    auto* inner = lua_torive<ScriptedPromise>(L, valueIdx, true);
    if (inner && inner != this)
    {
        if (inner->isFulfilled())
        {
            lua_rawgeti(L, LUA_REGISTRYINDEX, inner->resultRef());
            resolve(L, lua_gettop(L));
            lua_pop(L, 1);
            return;
        }
        if (inner->isRejected())
        {
            lua_rawgeti(L, LUA_REGISTRYINDEX, inner->resultRef());
            reject(L, lua_gettop(L));
            lua_pop(L, 1);
            return;
        }
        if (inner->isCancelled())
        {
            cancel(L);
            return;
        }
        // Still pending — wire this as a consumer of inner so that when
        // inner settles, it forwards the result to this promise via the
        // no-handler ThenCallback propagation path.
        // Release any existing parent ref (this promise may already be
        // chained from another promise via andThen).
        if (m_parentRef != LUA_NOREF)
            lua_unref(L, m_parentRef);
        lua_pushvalue(L, valueIdx); // push inner for ref
        m_parentRef = lua_ref(L, -1);
        lua_pop(L, 1);

        // selfRef is a registry ref to this promise's userdata — duplicate
        // it for the consumer list and the chained callback.
        lua_rawgeti(L, LUA_REGISTRYINDEX, selfRef);
        int consumerRef = lua_ref(L, -1); // for m_consumerRefs
        lua_pop(L, 1);

        lua_rawgeti(L, LUA_REGISTRYINDEX, selfRef);
        int chainedRef = lua_ref(L, -1); // for ThenCallback
        lua_pop(L, 1);

        inner->m_consumerRefs.push_back(consumerRef);
        {
            ScriptedPromise::ThenCallback cb;
            cb.chainedPromiseRef = chainedRef;
            inner->m_thenCallbacks.push_back(cb);
        }
        return;
    }

    m_promiseState = PromiseState::Fulfilled;
    lua_pushvalue(L, valueIdx);
    m_resultRef = lua_ref(L, -1);
    lua_pop(L, 1);
    promise_notifyCallbacks(L, this);
    promise_cleanupLinks(L, this);
}

void ScriptedPromise::reject(lua_State* L, int errorIdx)
{
    if (m_promiseState != PromiseState::Pending)
        return;
    m_promiseState = PromiseState::Rejected;
    lua_pushvalue(L, errorIdx);
    m_resultRef = lua_ref(L, -1);
    lua_pop(L, 1);
    promise_notifyCallbacks(L, this);
    promise_cleanupLinks(L, this);
}

void ScriptedPromise::cancel(lua_State* L)
{
    if (m_promiseState != PromiseState::Pending)
        return;

    m_promiseState = PromiseState::Cancelled;

    // 1. Fire the onCancel hook (if registered).
    if (m_onCancelRef != LUA_NOREF)
    {
        lua_rawgeti(L, LUA_REGISTRYINDEX, m_onCancelRef);
        if (lua_pcall(L, 0, 0, 0) != LUA_OK)
        {
            lua_pop(L, 1); // pop error object
        }
    }

    // 2. Propagate cancel DOWN to all consumers.
    for (int ref : m_consumerRefs)
    {
        if (ref != LUA_NOREF)
        {
            lua_rawgeti(L, LUA_REGISTRYINDEX, ref);
            auto* consumer = lua_torive<ScriptedPromise>(L, -1, true);
            lua_pop(L, 1);
            if (consumer && consumer->isPending())
                consumer->cancel(L);
        }
    }

    // 3. Notify parent: cancel parent if ALL its consumers are now cancelled.
    if (m_parentRef != LUA_NOREF)
    {
        lua_rawgeti(L, LUA_REGISTRYINDEX, m_parentRef);
        auto* parent = lua_torive<ScriptedPromise>(L, -1, true);
        lua_pop(L, 1);
        if (parent && parent->isPending())
        {
            bool allCancelled = true;
            for (int cRef : parent->m_consumerRefs)
            {
                if (cRef != LUA_NOREF)
                {
                    lua_rawgeti(L, LUA_REGISTRYINDEX, cRef);
                    auto* c = lua_torive<ScriptedPromise>(L, -1, true);
                    lua_pop(L, 1);
                    if (c && !c->isCancelled())
                    {
                        allCancelled = false;
                        break;
                    }
                }
            }
            if (allCancelled)
                parent->cancel(L);
        }
    }

    // 4. Clean up callback refs (they will never fire).
    for (auto& cb : m_thenCallbacks)
    {
        // Fire cancelRef if present (async/await cleanup).
        if (cb.cancelRef != LUA_NOREF)
        {
            lua_rawgeti(L, LUA_REGISTRYINDEX, cb.cancelRef);
            if (lua_pcall(L, 0, 0, 0) != LUA_OK)
            {
                lua_pop(L, 1);
            }
            lua_unref(L, cb.cancelRef);
        }
        if (cb.successRef != LUA_NOREF)
            lua_unref(L, cb.successRef);
        if (cb.failureRef != LUA_NOREF)
            lua_unref(L, cb.failureRef);
        if (cb.chainedPromiseRef != LUA_NOREF)
            lua_unref(L, cb.chainedPromiseRef);
    }
    m_thenCallbacks.clear();
    for (auto& cb : m_finallyCallbacks)
    {
        if (cb.callbackRef != LUA_NOREF)
            lua_unref(L, cb.callbackRef);
        if (cb.chainedPromiseRef != LUA_NOREF)
            lua_unref(L, cb.chainedPromiseRef);
    }
    m_finallyCallbacks.clear();

    // 5. Clean up parent/consumer/onCancel refs.
    promise_cleanupLinks(L, this);
}

ScriptedPromise::~ScriptedPromise()
{
    // Release any held Lua refs.  m_state may be the main thread of a
    // still-alive VM, so only unref if we actually have refs.
    if (m_resultRef != LUA_NOREF && m_state)
        lua_unref(m_state, m_resultRef);
    for (auto& cb : m_thenCallbacks)
    {
        if (cb.successRef != LUA_NOREF && m_state)
            lua_unref(m_state, cb.successRef);
        if (cb.failureRef != LUA_NOREF && m_state)
            lua_unref(m_state, cb.failureRef);
        if (cb.chainedPromiseRef != LUA_NOREF && m_state)
            lua_unref(m_state, cb.chainedPromiseRef);
        if (cb.cancelRef != LUA_NOREF && m_state)
            lua_unref(m_state, cb.cancelRef);
    }
    for (auto& cb : m_finallyCallbacks)
    {
        if (cb.callbackRef != LUA_NOREF && m_state)
            lua_unref(m_state, cb.callbackRef);
        if (cb.chainedPromiseRef != LUA_NOREF && m_state)
            lua_unref(m_state, cb.chainedPromiseRef);
    }
    if (m_parentRef != LUA_NOREF && m_state)
        lua_unref(m_state, m_parentRef);
    for (int ref : m_consumerRefs)
    {
        if (ref != LUA_NOREF && m_state)
            lua_unref(m_state, ref);
    }
    if (m_onCancelRef != LUA_NOREF && m_state)
        lua_unref(m_state, m_onCancelRef);
}

// ── callback dispatch ──────────────────────────────────────────────────────

static void promise_notifyCallbacks(lua_State* L, ScriptedPromise* p)
{
    // Cancelled promises don't fire andThen/catch/finally handlers.
    if (p->isCancelled())
        return;

    bool fulfilled = p->isFulfilled();

    // Push the result/error value once.
    lua_rawgeti(L, LUA_REGISTRYINDEX, p->resultRef());

    int resultIdx = lua_gettop(L);

    // ── andThen / catch callbacks ──
    for (auto& cb : p->m_thenCallbacks)
    {
        // Resolve the chained promise ref to a pointer.
        ScriptedPromise* chained = nullptr;
        if (cb.chainedPromiseRef != LUA_NOREF)
        {
            lua_rawgeti(L, LUA_REGISTRYINDEX, cb.chainedPromiseRef);
            chained = lua_torive<ScriptedPromise>(L, -1, true);
            lua_pop(L, 1);
            // Skip cancelled consumers.
            if (chained && chained->isCancelled())
                chained = nullptr;
        }

        int handlerRef = fulfilled ? cb.successRef : cb.failureRef;
        if (handlerRef != LUA_NOREF)
        {
            // Call the handler with the result.
            lua_rawgeti(L, LUA_REGISTRYINDEX, handlerRef);
            lua_pushvalue(L, resultIdx);
            int status = lua_pcall(L, 1, 1, 0);
            if (status == LUA_OK)
            {
                // Handler returned a value — resolve chained promise with it.
                // Use the selfRef overload for Promise/A+ flattening.
                if (chained)
                    chained->resolve(L, lua_gettop(L), cb.chainedPromiseRef);
                lua_pop(L, 1);
            }
            else
            {
                // Handler errored — reject chained promise.
                if (chained)
                    chained->reject(L, lua_gettop(L));
                lua_pop(L, 1); // pop error
            }
        }
        else
        {
            // No handler for this branch — propagate to chained promise.
            // No flattening needed here — the value is already unwrapped
            // from the parent, not a handler return value.
            if (chained)
            {
                if (fulfilled)
                    chained->resolve(L, resultIdx, cb.chainedPromiseRef);
                else
                    chained->reject(L, resultIdx);
            }
        }

        // Unref handlers.
        if (cb.successRef != LUA_NOREF)
            lua_unref(L, cb.successRef);
        if (cb.failureRef != LUA_NOREF)
            lua_unref(L, cb.failureRef);
        if (cb.chainedPromiseRef != LUA_NOREF)
            lua_unref(L, cb.chainedPromiseRef);
        if (cb.cancelRef != LUA_NOREF)
            lua_unref(L, cb.cancelRef);
    }

    // ── finally callbacks ──
    for (auto& cb : p->m_finallyCallbacks)
    {
        ScriptedPromise* chained = nullptr;
        if (cb.chainedPromiseRef != LUA_NOREF)
        {
            lua_rawgeti(L, LUA_REGISTRYINDEX, cb.chainedPromiseRef);
            chained = lua_torive<ScriptedPromise>(L, -1, true);
            lua_pop(L, 1);
            if (chained && chained->isCancelled())
                chained = nullptr;
        }

        if (cb.callbackRef != LUA_NOREF)
        {
            lua_rawgeti(L, LUA_REGISTRYINDEX, cb.callbackRef);
            int status = lua_pcall(L, 0, 0, 0);
            if (status != LUA_OK)
            {
                // finally threw — reject the chained promise with the error
                // instead of propagating the original result.
                if (chained)
                    chained->reject(L, lua_gettop(L));
                lua_pop(L, 1); // pop error from stack
                lua_unref(L, cb.callbackRef);
                if (cb.chainedPromiseRef != LUA_NOREF)
                    lua_unref(L, cb.chainedPromiseRef);
                continue;
            }
            lua_unref(L, cb.callbackRef);
        }
        // Propagate original result to chained promise.
        if (chained)
        {
            if (fulfilled)
                chained->resolve(L, resultIdx);
            else
                chained->reject(L, resultIdx);
        }
        if (cb.chainedPromiseRef != LUA_NOREF)
            lua_unref(L, cb.chainedPromiseRef);
    }

    lua_pop(L, 1); // pop result value

    p->m_thenCallbacks.clear();
    p->m_finallyCallbacks.clear();
}

// ── helper: wire parent/consumer links for a chained promise ────────────

// Call after creating a chained promise. `source` is at stack index
// `sourceIdx`, chained promise is at the top of the stack.
static void promise_wireLinks(lua_State* L,
                              ScriptedPromise* source,
                              int sourceIdx,
                              ScriptedPromise* chained)
{
    // Consumer ref: source → chained (for downward cancel propagation).
    lua_pushvalue(L, -1); // push chained
    int consumerRef = lua_ref(L, -1);
    lua_pop(L, 1);
    source->m_consumerRefs.push_back(consumerRef);

    // Parent ref: chained → source (for upward cancel propagation).
    lua_pushvalue(L, sourceIdx);
    chained->m_parentRef = lua_ref(L, -1);
    lua_pop(L, 1);
}

// ── namecall: p:andThen(), p:catch(), p:finally(), p:cancel(), etc. ─────

static int promise_namecall(lua_State* L)
{
    int atom;
    const char* str = lua_namecallatom(L, &atom);
    auto* p = lua_torive<ScriptedPromise>(L, 1);

    if (atom == (int)LuaAtoms::andThen)
    {
        // andThen(onFulfilled [, onRejected]) -> Promise
        int successRef = LUA_NOREF;
        int failureRef = LUA_NOREF;
        if (lua_isfunction(L, 2))
        {
            lua_pushvalue(L, 2);
            successRef = lua_ref(L, -1);
            lua_pop(L, 1);
        }
        if (lua_isfunction(L, 3))
        {
            lua_pushvalue(L, 3);
            failureRef = lua_ref(L, -1);
            lua_pop(L, 1);
        }

        auto* chained = lua_newrive<ScriptedPromise>(L, lua_mainthread(L));
        // The chained promise is now at the top of the stack.
        // Create a registry ref to prevent GC while the parent is pending.
        lua_pushvalue(L, -1);
        int chainedRef = lua_ref(L, -1);
        lua_pop(L, 1);

        // Wire parent/consumer links.
        promise_wireLinks(L, p, 1, chained);

        {
            ScriptedPromise::ThenCallback cb;
            cb.successRef = successRef;
            cb.failureRef = failureRef;
            cb.chainedPromiseRef = chainedRef;
            p->m_thenCallbacks.push_back(cb);
        }

        if (!p->isPending())
            promise_notifyCallbacks(L, p);

        return 1; // return chained promise (still on stack)
    }
    else if (atom == (int)LuaAtoms::catch_)
    {
        // catch(onRejected) -> Promise  (sugar for andThen(nil, onRejected))
        int failureRef = LUA_NOREF;
        if (lua_isfunction(L, 2))
        {
            lua_pushvalue(L, 2);
            failureRef = lua_ref(L, -1);
            lua_pop(L, 1);
        }

        auto* chained = lua_newrive<ScriptedPromise>(L, lua_mainthread(L));
        lua_pushvalue(L, -1);
        int chainedRef = lua_ref(L, -1);
        lua_pop(L, 1);

        promise_wireLinks(L, p, 1, chained);

        {
            ScriptedPromise::ThenCallback cb;
            cb.failureRef = failureRef;
            cb.chainedPromiseRef = chainedRef;
            p->m_thenCallbacks.push_back(cb);
        }

        if (!p->isPending())
            promise_notifyCallbacks(L, p);

        return 1;
    }
    else if (atom == (int)LuaAtoms::finally_)
    {
        int cbRef = LUA_NOREF;
        if (lua_isfunction(L, 2))
        {
            lua_pushvalue(L, 2);
            cbRef = lua_ref(L, -1);
            lua_pop(L, 1);
        }

        auto* chained = lua_newrive<ScriptedPromise>(L, lua_mainthread(L));
        lua_pushvalue(L, -1);
        int chainedRef = lua_ref(L, -1);
        lua_pop(L, 1);

        promise_wireLinks(L, p, 1, chained);

        {
            ScriptedPromise::FinallyCallback cb;
            cb.callbackRef = cbRef;
            cb.chainedPromiseRef = chainedRef;
            p->m_finallyCallbacks.push_back(cb);
        }

        if (!p->isPending())
            promise_notifyCallbacks(L, p);

        return 1;
    }
    else if (atom == (int)LuaAtoms::cancel)
    {
        p->cancel(L);
        lua_pushvalue(L, 1); // return self for chaining
        return 1;
    }
    else if (atom == (int)LuaAtoms::onCancel)
    {
        // onCancel(fn) — register a cancellation hook.
        if (lua_isfunction(L, 2))
        {
            if (p->m_onCancelRef != LUA_NOREF)
                lua_unref(L, p->m_onCancelRef);
            lua_pushvalue(L, 2);
            p->m_onCancelRef = lua_ref(L, -1);
            lua_pop(L, 1);
        }
        lua_pushvalue(L, 1); // return self
        return 1;
    }
    else if (atom == (int)LuaAtoms::getStatus)
    {
        switch (p->m_promiseState)
        {
            case PromiseState::Pending:
                lua_pushstring(L, "Pending");
                break;
            case PromiseState::Fulfilled:
                lua_pushstring(L, "Fulfilled");
                break;
            case PromiseState::Rejected:
                lua_pushstring(L, "Rejected");
                break;
            case PromiseState::Cancelled:
                lua_pushstring(L, "Cancelled");
                break;
        }
        return 1;
    }

    luaL_error(L, "'%s' is not a valid method of Promise", str);
    return 0;
}

// ── static methods: Promise.resolve(), Promise.reject(), Promise.all() ───

static int promise_static_resolve(lua_State* L)
{
    auto* p = lua_newrive<ScriptedPromise>(L, lua_mainthread(L));
    int pIdx = lua_gettop(L); // p is on top of stack
    if (pIdx >= 2)
    {
        // Promise.resolve(value) — use selfRef overload for A+ flattening.
        lua_pushvalue(L, pIdx);
        int selfRef = lua_ref(L, -1);
        lua_pop(L, 1);
        p->resolve(L, 1, selfRef); // arg 1 is the value (before we pushed p)
        lua_unref(L, selfRef);
    }
    else
    {
        // Promise.resolve() with no value → resolve with nil
        lua_pushnil(L);
        p->resolve(L, lua_gettop(L));
        lua_pop(L, 1);
    }
    return 1;
}

static int promise_static_reject(lua_State* L)
{
    auto* p = lua_newrive<ScriptedPromise>(L, lua_mainthread(L));
    if (lua_gettop(L) >= 2)
    {
        p->reject(L, 1);
    }
    else
    {
        lua_pushstring(L, "rejected");
        p->reject(L, lua_gettop(L));
        lua_pop(L, 1);
    }
    return 1;
}

// ── Promise.new(executor) ───────────────────────────────────────────────

static int promise_static_new(lua_State* L)
{
    luaL_checktype(L, 1, LUA_TFUNCTION);

    auto* p = lua_newrive<ScriptedPromise>(L, lua_mainthread(L));
    int promiseIdx = lua_gettop(L);

    // Create resolve callback — captures promise as direct upvalue (no
    // registry ref needed; GC keeps it alive via the closure).
    lua_pushvalue(L, promiseIdx);
    lua_pushcclosurek(
        L,
        [](lua_State* L) -> int {
            auto* promise =
                lua_torive<ScriptedPromise>(L, lua_upvalueindex(1), true);
            if (promise && promise->isPending())
            {
                if (lua_gettop(L) >= 1)
                {
                    // Use selfRef overload for A+ flattening.
                    lua_pushvalue(L, lua_upvalueindex(1));
                    int selfRef = lua_ref(L, -1);
                    lua_pop(L, 1);
                    promise->resolve(L, 1, selfRef);
                    lua_unref(L, selfRef);
                }
                else
                {
                    lua_pushnil(L);
                    promise->resolve(L, lua_gettop(L));
                    lua_pop(L, 1);
                }
            }
            return 0;
        },
        nullptr,
        1,
        nullptr);
    int resolveIdx = lua_gettop(L);

    // Create reject callback.
    lua_pushvalue(L, promiseIdx);
    lua_pushcclosurek(
        L,
        [](lua_State* L) -> int {
            auto* promise =
                lua_torive<ScriptedPromise>(L, lua_upvalueindex(1), true);
            if (promise && promise->isPending())
            {
                if (lua_gettop(L) >= 1)
                    promise->reject(L, 1);
                else
                {
                    lua_pushstring(L, "rejected");
                    promise->reject(L, lua_gettop(L));
                    lua_pop(L, 1);
                }
            }
            return 0;
        },
        nullptr,
        1,
        nullptr);
    int rejectIdx = lua_gettop(L);

    // Create onCancel callback.
    lua_pushvalue(L, promiseIdx);
    lua_pushcclosurek(
        L,
        [](lua_State* L) -> int {
            auto* promise =
                lua_torive<ScriptedPromise>(L, lua_upvalueindex(1), true);
            if (promise && lua_isfunction(L, 1))
            {
                if (promise->m_onCancelRef != LUA_NOREF)
                    lua_unref(L, promise->m_onCancelRef);
                lua_pushvalue(L, 1);
                promise->m_onCancelRef = lua_ref(L, -1);
                lua_pop(L, 1);
            }
            return 0;
        },
        nullptr,
        1,
        nullptr);
    int onCancelIdx = lua_gettop(L);

    // Call executor(resolve, reject, onCancel).
    lua_pushvalue(L, 1); // push executor
    lua_pushvalue(L, resolveIdx);
    lua_pushvalue(L, rejectIdx);
    lua_pushvalue(L, onCancelIdx);
    int status = lua_pcall(L, 3, 0, 0);
    if (status != LUA_OK)
    {
        // Executor threw — reject the promise with the error.
        if (p->isPending())
            p->reject(L, lua_gettop(L));
        lua_pop(L, 1); // pop error
    }

    // Clean up intermediates and return the promise.
    lua_pushvalue(L, promiseIdx);
    return 1;
}

// ── Promise.all ─────────────────────────────────────────────────────────

static int promise_static_all(lua_State* L)
{
    // Promise.all({p1, p2, ...}) -> Promise<{v1, v2, ...}>
    luaL_checktype(L, 1, LUA_TTABLE);
    int n = lua_objlen(L, 1);

    auto* result = lua_newrive<ScriptedPromise>(L, lua_mainthread(L));
    int resultIdx = lua_gettop(L);

    if (n == 0)
    {
        // Empty array → resolve immediately with empty table.
        lua_newtable(L);
        result->resolve(L, lua_gettop(L));
        lua_pop(L, 1);
        lua_pushvalue(L, resultIdx);
        return 1;
    }

    // Shared state table — stays on L's stack, captured as direct upvalue
    // by all success/failure closures. No registry ref needed.
    // Fields: { remaining=n, results={}, done=false }
    lua_newtable(L);
    int stateIdx = lua_gettop(L);
    lua_pushinteger(L, n);
    lua_setfield(L, stateIdx, "remaining");
    lua_newtable(L);
    lua_setfield(L, stateIdx, "results");
    lua_pushboolean(L, false);
    lua_setfield(L, stateIdx, "done");

    // Chained refs table — stays on stack, captured as direct upvalue.
    // Populated during the loop below with registry refs to chained
    // promises for cancellation cleanup.
    lua_newtable(L);
    int chainedRefsTableIdx = lua_gettop(L);

    for (int i = 1; i <= n; i++)
    {
        lua_rawgeti(L, 1, i); // push promises[i]
        auto* pi = lua_torive<ScriptedPromise>(L, -1, true);
        int piIdx = lua_gettop(L);
        if (!pi)
        {
            lua_pop(L, 1);
            luaL_error(L, "Promise.all: element %d is not a Promise", i);
            return 0;
        }

        // Create chained promise for this input.
        auto* chained = lua_newrive<ScriptedPromise>(L, lua_mainthread(L));
        int chainedIdx = lua_gettop(L);

        // Wire parent/consumer links: pi → chained.
        promise_wireLinks(L, pi, piIdx, chained);

        // Registry ref for ThenCallback dispatch.
        lua_pushvalue(L, chainedIdx);
        int chainedRef = lua_ref(L, -1);
        lua_pop(L, 1);

        // Store a separate ref in the chainedRefs table for cancellation.
        lua_pushvalue(L, chainedIdx);
        int cancelCollectionRef = lua_ref(L, -1);
        lua_pop(L, 1);
        lua_pushinteger(L, cancelCollectionRef);
        lua_rawseti(L, chainedRefsTableIdx, i);

        // Success closure: upvalues = (index, state, resultPromise,
        // chainedRefs)
        lua_pushinteger(L, i);
        lua_pushvalue(L, stateIdx);
        lua_pushvalue(L, resultIdx);
        lua_pushvalue(L, chainedRefsTableIdx);
        lua_pushcclosurek(
            L,
            [](lua_State* L) -> int {
                int idx = (int)lua_tointeger(L, lua_upvalueindex(1));
                // upvalue 2 = state table, 3 = result promise, 4 = chainedRefs

                lua_pushvalue(L, lua_upvalueindex(2)); // state
                int st = lua_gettop(L);

                lua_getfield(L, st, "done");
                if (lua_toboolean(L, -1))
                {
                    lua_pop(L, 2);
                    return 0;
                }
                lua_pop(L, 1);

                // Store result at index.
                lua_getfield(L, st, "results");
                lua_pushvalue(L, 1);
                lua_rawseti(L, -2, idx);
                lua_pop(L, 1);

                // Decrement remaining.
                lua_getfield(L, st, "remaining");
                int rem = (int)lua_tointeger(L, -1) - 1;
                lua_pop(L, 1);
                lua_pushinteger(L, rem);
                lua_setfield(L, st, "remaining");

                if (rem == 0)
                {
                    lua_pushboolean(L, true);
                    lua_setfield(L, st, "done");

                    auto* rp = lua_torive<ScriptedPromise>(L,
                                                           lua_upvalueindex(3),
                                                           true);
                    if (rp && rp->isPending())
                    {
                        lua_getfield(L, st, "results");
                        rp->resolve(L, lua_gettop(L));
                        lua_pop(L, 1);
                    }

                    // Clean up cancelCollectionRefs in chainedRefs table.
                    lua_pushvalue(L, lua_upvalueindex(4));
                    int tbl = lua_gettop(L);
                    int count = (int)lua_objlen(L, tbl);
                    for (int j = 1; j <= count; j++)
                    {
                        lua_rawgeti(L, tbl, j);
                        int cRef = (int)lua_tointeger(L, -1);
                        lua_pop(L, 1);
                        if (cRef != LUA_NOREF)
                            lua_unref(L, cRef);
                    }
                    lua_pop(L, 1); // pop table
                }

                lua_pop(L, 1); // pop state table
                return 0;
            },
            nullptr,
            4,
            nullptr);
        int successRef = lua_ref(L, -1);
        lua_pop(L, 1);

        // Failure closure: upvalues = (state, resultPromise, chainedRefs)
        lua_pushvalue(L, stateIdx);
        lua_pushvalue(L, resultIdx);
        lua_pushvalue(L, chainedRefsTableIdx);
        lua_pushcclosurek(
            L,
            [](lua_State* L) -> int {
                // upvalue 1 = state, 2 = result promise, 3 = chainedRefs
                lua_pushvalue(L, lua_upvalueindex(1)); // state
                int st = lua_gettop(L);

                lua_getfield(L, st, "done");
                if (lua_toboolean(L, -1))
                {
                    lua_pop(L, 2);
                    return 0;
                }
                lua_pop(L, 1);

                lua_pushboolean(L, true);
                lua_setfield(L, st, "done");

                auto* rp =
                    lua_torive<ScriptedPromise>(L, lua_upvalueindex(2), true);
                if (rp && rp->isPending())
                {
                    rp->reject(L, 1); // arg 1 is the error
                }

                // Cancel all chained promises from chainedRefs table.
                lua_pushvalue(L, lua_upvalueindex(3));
                int tbl = lua_gettop(L);
                int count = (int)lua_objlen(L, tbl);
                for (int j = 1; j <= count; j++)
                {
                    lua_rawgeti(L, tbl, j);
                    int cRef = (int)lua_tointeger(L, -1);
                    lua_pop(L, 1);
                    if (cRef != LUA_NOREF)
                    {
                        lua_rawgeti(L, LUA_REGISTRYINDEX, cRef);
                        auto* cp = lua_torive<ScriptedPromise>(L, -1, true);
                        lua_pop(L, 1);
                        if (cp && cp->isPending())
                            cp->cancel(L);
                        lua_unref(L, cRef);
                    }
                }
                lua_pop(L, 1); // pop table

                lua_pop(L, 1); // pop state table
                return 0;
            },
            nullptr,
            3,
            nullptr);
        int failureRef = lua_ref(L, -1);
        lua_pop(L, 1);

        {
            ScriptedPromise::ThenCallback cb;
            cb.successRef = successRef;
            cb.failureRef = failureRef;
            cb.chainedPromiseRef = chainedRef;
            pi->m_thenCallbacks.push_back(cb);
        }

        if (!pi->isPending())
            promise_notifyCallbacks(L, pi);

        lua_pop(L, 2); // pop chained promise and pi
    }

    // Register onCancel hook: captures chainedRefs table as direct upvalue.
    lua_pushvalue(L, chainedRefsTableIdx);
    lua_pushcclosurek(
        L,
        [](lua_State* L) -> int {
            // upvalue 1 = chainedRefs table
            lua_pushvalue(L, lua_upvalueindex(1));
            int tbl = lua_gettop(L);
            int count = (int)lua_objlen(L, tbl);

            for (int j = 1; j <= count; j++)
            {
                lua_rawgeti(L, tbl, j);
                int cRef = (int)lua_tointeger(L, -1);
                lua_pop(L, 1);
                if (cRef != LUA_NOREF)
                {
                    lua_rawgeti(L, LUA_REGISTRYINDEX, cRef);
                    auto* cp = lua_torive<ScriptedPromise>(L, -1, true);
                    lua_pop(L, 1);
                    if (cp && cp->isPending())
                        cp->cancel(L);
                    lua_unref(L, cRef);
                }
            }
            lua_pop(L, 1); // pop table
            return 0;
        },
        nullptr,
        1,
        nullptr);
    result->m_onCancelRef = lua_ref(L, -1);
    lua_pop(L, 1);

    // Pop the chainedRefs table and state table from the stack.
    lua_pop(L, 2);

    lua_pushvalue(L, resultIdx);
    return 1;
}

// ── Promise table __index (static method dispatch) ─────────────────────

static int promise_table_index(lua_State* L)
{
    const char* key = luaL_checkstring(L, 2);
    if (strcmp(key, "resolve") == 0)
    {
        lua_pushcfunction(L, promise_static_resolve, "Promise.resolve");
        return 1;
    }
    if (strcmp(key, "reject") == 0)
    {
        lua_pushcfunction(L, promise_static_reject, "Promise.reject");
        return 1;
    }
    if (strcmp(key, "all") == 0)
    {
        lua_pushcfunction(L, promise_static_all, "Promise.all");
        return 1;
    }
    if (strcmp(key, "new") == 0)
    {
        lua_pushcfunction(L, promise_static_new, "Promise.new");
        return 1;
    }
    luaL_error(L, "'%s' is not a valid member of Promise", key);
    return 0;
}

// ── async / await ──────────────────────────────────────────────────────────

// await(promise) -> (ok: boolean, value: any)
// Non-throwing: returns (true, resolvedValue) or (false, rejectionReason).
// Matches Roblox Promise:await() semantics. Use expect() for throwing variant.
static int lua_await(lua_State* L)
{
    auto* p = lua_torive<ScriptedPromise>(L, 1);

    if (!lua_isyieldable(L))
    {
        luaL_error(L, "await() must be called inside async()");
        return 0;
    }

    if (p->isFulfilled())
    {
        lua_pushboolean(L, true);
        lua_rawgeti(L, LUA_REGISTRYINDEX, p->resultRef());
        return 2;
    }
    if (p->isRejected())
    {
        lua_pushboolean(L, false);
        lua_rawgeti(L, LUA_REGISTRYINDEX, p->resultRef());
        return 2;
    }
    if (p->isCancelled())
    {
        lua_pushboolean(L, false);
        lua_pushstring(L, "Promise was cancelled");
        return 2;
    }

    // Pending — yield with the promise on the stack so the caller
    // (async's resume callback) can attach andThen to it.
    lua_pushvalue(L, 1);
    int nresults = lua_yield(L, 1);
    // After resume: success callback pushes (true, value),
    // failure callback pushes (false, error).
    // Return them directly — no throwing.
    return nresults;
}

static int lua_async(lua_State* L)
{
    luaL_checktype(L, 1, LUA_TFUNCTION);

    // Create the result promise.
    lua_newrive<ScriptedPromise>(L, lua_mainthread(L));
    int promiseIdx = lua_gettop(L);

    // Create a coroutine. Both stay on L's stack — closures capture
    // them as direct upvalues, GC keeps them alive.
    lua_State* co = lua_newthread(L);
    lua_setthreaddata(co, lua_getthreaddata(L));
    int coIdx = lua_gettop(L);

    // Copy function to coroutine stack.
    lua_pushvalue(L, 1);
    lua_xmove(L, co, 1);

    // Resume the coroutine (0 args).
    int status = lua_resume(co, L, 0);

    handleCoroutineCompletion(L, co, status, coIdx, promiseIdx);

    // Clean up thread from main stack.
    lua_remove(L, coIdx);

    // Return the promise.
    lua_pushvalue(L, promiseIdx);
    return 1;
}

// handleCoroutineCompletion takes the coroutine and async promise as
// stack values (coIdx, asyncPromiseIdx on L). Closures capture them
// as direct upvalues — no registry refs needed; GC owns lifetime.
static void handleCoroutineCompletion(lua_State* L,
                                      lua_State* co,
                                      int status,
                                      int coIdx,
                                      int asyncPromiseIdx)
{
    auto* asyncPromise = lua_torive<ScriptedPromise>(L, asyncPromiseIdx, true);

    if (status == LUA_OK)
    {
        // Coroutine finished — resolve with the return value (if any).
        // Use selfRef overload for A+ flattening (user may return a Promise).
        if (lua_gettop(co) > 0)
        {
            lua_xmove(co, L, 1);
            lua_pushvalue(L, asyncPromiseIdx);
            int selfRef = lua_ref(L, -1);
            lua_pop(L, 1);
            asyncPromise->resolve(L, lua_gettop(L), selfRef);
            lua_unref(L, selfRef);
            lua_pop(L, 1);
        }
        else
        {
            lua_pushnil(L);
            asyncPromise->resolve(L, lua_gettop(L));
            lua_pop(L, 1);
        }
        // No refs to clean up — GC handles the upvalues.
    }
    else if (status == LUA_YIELD)
    {
        // Coroutine yielded at await(promise).
        if (lua_gettop(co) < 1)
        {
            luaL_error(L, "async: coroutine yielded without a promise");
            return;
        }

        auto* awaited = lua_torive<ScriptedPromise>(co, -1, true);
        if (!awaited)
        {
            luaL_error(L, "async: await() argument is not a Promise");
            return;
        }
        lua_pop(co, 1); // pop the promise from co stack

        // Success callback: captures coroutine + asyncPromise as upvalues.
        // asyncPromise is captured to prevent GC, not used directly.
        lua_pushvalue(L, coIdx);
        lua_pushvalue(L, asyncPromiseIdx);
        lua_pushcclosurek(
            L,
            [](lua_State* L) -> int {
                lua_State* co = lua_tothread(L, lua_upvalueindex(1));

                // Push (true, resolvedValue) onto coroutine stack.
                lua_pushboolean(co, true);
                lua_pushvalue(L, 1);
                lua_xmove(L, co, 1);

                int status = lua_resume(co, L, 2);

                // Push co + asyncPromise onto stack for recursive call.
                lua_pushvalue(L, lua_upvalueindex(1));
                int coIdx = lua_gettop(L);
                lua_pushvalue(L, lua_upvalueindex(2));
                int apIdx = lua_gettop(L);
                handleCoroutineCompletion(L, co, status, coIdx, apIdx);
                lua_pop(L, 2);
                return 0;
            },
            nullptr,
            2,
            nullptr);
        int successRef = lua_ref(L, -1);
        lua_pop(L, 1);

        // Failure callback: resume coroutine with (false, error).
        // asyncPromise is captured to prevent GC, not used directly.
        lua_pushvalue(L, coIdx);
        lua_pushvalue(L, asyncPromiseIdx);
        lua_pushcclosurek(
            L,
            [](lua_State* L) -> int {
                lua_State* co = lua_tothread(L, lua_upvalueindex(1));

                lua_pushboolean(co, false);
                lua_pushvalue(L, 1);
                lua_xmove(L, co, 1);

                int status = lua_resume(co, L, 2);

                lua_pushvalue(L, lua_upvalueindex(1));
                int coIdx = lua_gettop(L);
                lua_pushvalue(L, lua_upvalueindex(2));
                int apIdx = lua_gettop(L);
                handleCoroutineCompletion(L, co, status, coIdx, apIdx);
                lua_pop(L, 2);
                return 0;
            },
            nullptr,
            2,
            nullptr);
        int failureRef = lua_ref(L, -1);
        lua_pop(L, 1);

        // Cancel callback: close the coroutine and reject async promise.
        lua_pushvalue(L, coIdx);
        lua_pushvalue(L, asyncPromiseIdx);
        lua_pushcclosurek(
            L,
            [](lua_State* L) -> int {
                lua_State* co = lua_tothread(L, lua_upvalueindex(1));
                auto* ap =
                    lua_torive<ScriptedPromise>(L, lua_upvalueindex(2), true);

                // Close the coroutine to release its stack/locals
                // immediately rather than waiting for GC.
                if (co)
                    lua_resetthread(co);

                if (ap && ap->isPending())
                {
                    lua_pushstring(L, "Promise was cancelled");
                    ap->reject(L, lua_gettop(L));
                    lua_pop(L, 1);
                }
                return 0;
            },
            nullptr,
            2,
            nullptr);
        int cancelRef = lua_ref(L, -1);
        lua_pop(L, 1);

        {
            ScriptedPromise::ThenCallback cb;
            cb.successRef = successRef;
            cb.failureRef = failureRef;
            cb.cancelRef = cancelRef;
            awaited->m_thenCallbacks.push_back(cb);
        }

        if (!awaited->isPending())
            promise_notifyCallbacks(L, awaited);
    }
    else
    {
        // Coroutine errored.
        if (lua_gettop(co) > 0)
        {
            lua_xmove(co, L, 1);
            asyncPromise->reject(L, lua_gettop(L));
            lua_pop(L, 1);
        }
        else
        {
            lua_pushstring(L, "async function errored");
            asyncPromise->reject(L, lua_gettop(L));
            lua_pop(L, 1);
        }
    }
}

// ── registration ───────────────────────────────────────────────────────────

int luaopen_rive_promise(lua_State* L)
{
    // Register ScriptedPromise type.
    lua_register_rive<ScriptedPromise>(L);
    lua_pushcfunction(L, promise_namecall, nullptr);
    lua_setfield(L, -2, "__namecall");
    lua_setreadonly(L, -1, true);
    lua_pop(L, 1); // pop metatable

    // Create global Promise table.
    lua_newtable(L);

    // Metatable for static method dispatch.
    lua_newtable(L);
    lua_pushcfunction(L, promise_table_index, nullptr);
    lua_setfield(L, -2, "__index");
    lua_setmetatable(L, -2);

    lua_setglobal(L, "Promise");

    // Register global await / async functions.
    lua_pushcfunction(L, lua_await, "await");
    lua_setglobal(L, "await");

    lua_pushcfunction(L, lua_async, "async");
    lua_setglobal(L, "async");

    return 0;
}

} // namespace rive

#endif // WITH_RIVE_SCRIPTING