// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
//
// SPDX-License-Identifier: Apache-2.0

#include <Python.h>

#include "resource_handles.hpp"
#include <cuda.h>
#include <atomic>
#include <array>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <list>
#include <map>
#include <mutex>
#include <stdexcept>
#include <unordered_map>
#include <vector>

#ifndef _WIN32
#include <unistd.h>
#endif

namespace cuda_core {

// ============================================================================
// CUDA driver function pointers
//
// These are populated by _resource_handles.pyx at module import time using
// function pointers extracted from cuda.bindings.cydriver.__pyx_capi__.
// ============================================================================

decltype(&cuDevicePrimaryCtxRetain) p_cuDevicePrimaryCtxRetain = nullptr;
decltype(&cuDevicePrimaryCtxRelease) p_cuDevicePrimaryCtxRelease = nullptr;
decltype(&cuCtxGetCurrent) p_cuCtxGetCurrent = nullptr;
decltype(&cuGreenCtxCreate) p_cuGreenCtxCreate = nullptr;
decltype(&cuGreenCtxDestroy) p_cuGreenCtxDestroy = nullptr;
decltype(&cuCtxFromGreenCtx) p_cuCtxFromGreenCtx = nullptr;
decltype(&cuDevResourceGenerateDesc) p_cuDevResourceGenerateDesc = nullptr;

decltype(&cuGreenCtxStreamCreate) p_cuGreenCtxStreamCreate = nullptr;

decltype(&cuStreamCreateWithPriority) p_cuStreamCreateWithPriority = nullptr;
decltype(&cuStreamDestroy) p_cuStreamDestroy = nullptr;

decltype(&cuEventCreate) p_cuEventCreate = nullptr;
decltype(&cuEventDestroy) p_cuEventDestroy = nullptr;
decltype(&cuIpcOpenEventHandle) p_cuIpcOpenEventHandle = nullptr;

decltype(&cuDeviceGetCount) p_cuDeviceGetCount = nullptr;

decltype(&cuMemPoolSetAccess) p_cuMemPoolSetAccess = nullptr;
decltype(&cuMemPoolDestroy) p_cuMemPoolDestroy = nullptr;
decltype(&cuMemPoolCreate) p_cuMemPoolCreate = nullptr;
decltype(&cuDeviceGetMemPool) p_cuDeviceGetMemPool = nullptr;
decltype(&cuMemPoolImportFromShareableHandle) p_cuMemPoolImportFromShareableHandle = nullptr;

decltype(&cuMemAllocFromPoolAsync) p_cuMemAllocFromPoolAsync = nullptr;
decltype(&cuMemAllocAsync) p_cuMemAllocAsync = nullptr;
decltype(&cuMemAlloc) p_cuMemAlloc = nullptr;
decltype(&cuMemAllocHost) p_cuMemAllocHost = nullptr;

decltype(&cuMemFreeAsync) p_cuMemFreeAsync = nullptr;
decltype(&cuMemFree) p_cuMemFree = nullptr;
decltype(&cuMemFreeHost) p_cuMemFreeHost = nullptr;

decltype(&cuMemPoolImportPointer) p_cuMemPoolImportPointer = nullptr;

decltype(&cuLibraryLoadFromFile) p_cuLibraryLoadFromFile = nullptr;
decltype(&cuLibraryLoadData) p_cuLibraryLoadData = nullptr;
decltype(&cuLibraryUnload) p_cuLibraryUnload = nullptr;
decltype(&cuLibraryGetKernel) p_cuLibraryGetKernel = nullptr;

// Graph
decltype(&cuGraphDestroy) p_cuGraphDestroy = nullptr;
decltype(&cuGraphInstantiateWithParams) p_cuGraphInstantiateWithParams = nullptr;
decltype(&cuGraphExecUpdate) p_cuGraphExecUpdate = nullptr;
decltype(&cuGraphExecDestroy) p_cuGraphExecDestroy = nullptr;
decltype(&cuUserObjectCreate) p_cuUserObjectCreate = nullptr;
decltype(&cuUserObjectRelease) p_cuUserObjectRelease = nullptr;
decltype(&cuGraphRetainUserObject) p_cuGraphRetainUserObject = nullptr;
decltype(&cuGraphReleaseUserObject) p_cuGraphReleaseUserObject = nullptr;
decltype(&cuGraphNodeFindInClone) p_cuGraphNodeFindInClone = nullptr;
decltype(&cuGraphChildGraphNodeGetGraph) p_cuGraphChildGraphNodeGetGraph = nullptr;

// Linker
decltype(&cuLinkDestroy) p_cuLinkDestroy = nullptr;

// GL interop pointers
decltype(&cuGraphicsUnmapResources) p_cuGraphicsUnmapResources = nullptr;
decltype(&cuGraphicsUnregisterResource) p_cuGraphicsUnregisterResource = nullptr;

decltype(&cuArray3DCreate) p_cuArray3DCreate = nullptr;
decltype(&cuArrayDestroy) p_cuArrayDestroy = nullptr;
decltype(&cuMipmappedArrayCreate) p_cuMipmappedArrayCreate = nullptr;
decltype(&cuMipmappedArrayDestroy) p_cuMipmappedArrayDestroy = nullptr;
decltype(&cuMipmappedArrayGetLevel) p_cuMipmappedArrayGetLevel = nullptr;
decltype(&cuTexObjectCreate) p_cuTexObjectCreate = nullptr;
decltype(&cuTexObjectDestroy) p_cuTexObjectDestroy = nullptr;
decltype(&cuSurfObjectCreate) p_cuSurfObjectCreate = nullptr;
decltype(&cuSurfObjectDestroy) p_cuSurfObjectDestroy = nullptr;

// SM resource split (13.1+ — may be null on older drivers/bindings)
#if CUDA_VERSION >= 13010
decltype(&cuDevSmResourceSplit) p_cuDevSmResourceSplit = nullptr;
#else
void* p_cuDevSmResourceSplit = nullptr;
#endif

// NVRTC function pointers
decltype(&nvrtcDestroyProgram) p_nvrtcDestroyProgram = nullptr;

// NVVM function pointers (may be null if NVVM is not available)
NvvmDestroyProgramFn p_nvvmDestroyProgram = nullptr;

// nvJitLink function pointers (may be null if nvJitLink is not available)
NvJitLinkDestroyFn p_nvJitLinkDestroy = nullptr;

// ============================================================================
// GIL management helpers
// ============================================================================

namespace {

// Helper to release the GIL while calling into the CUDA driver.
// This guard is *conditional*: if the caller already dropped the GIL,
// we avoid calling PyEval_SaveThread (which requires holding the GIL).
// It also handles the case where Python is finalizing and GIL operations
// are no longer safe.
class GILReleaseGuard {
public:
    GILReleaseGuard() : tstate_(nullptr), released_(false) {
        // Don't try to manipulate GIL if Python is finalizing
        if (!Py_IsInitialized() || py_is_finalizing()) {
            return;
        }
        // PyGILState_Check() returns 1 if the GIL is held by this thread.
        if (PyGILState_Check()) {
            tstate_ = PyEval_SaveThread();
            released_ = true;
        }
        // Note: If the GIL is not released (finalizing, or not held):
        // - Reduces parallelism (other Python threads remain blocked)
        // - No deadlock risk as long as the guarded code doesn't call back into Python
    }

    ~GILReleaseGuard() {
        if (released_) {
            PyEval_RestoreThread(tstate_);
        }
    }

    // Non-copyable, non-movable
    GILReleaseGuard(const GILReleaseGuard&) = delete;
    GILReleaseGuard& operator=(const GILReleaseGuard&) = delete;

private:
    PyThreadState* tstate_;
    bool released_;
};

// Helper to acquire the GIL when we might not hold it.
// Use in C++ destructors that need to manipulate Python objects.
class GILAcquireGuard {
public:
    GILAcquireGuard() : acquired_(false) {
        // Don't try to acquire GIL if Python is finalizing
        if (!Py_IsInitialized() || py_is_finalizing()) {
            return;
        }
        gstate_ = PyGILState_Ensure();
        acquired_ = true;
    }

    ~GILAcquireGuard() {
        if (acquired_) {
            PyGILState_Release(gstate_);
        }
    }

    bool acquired() const { return acquired_; }

    // Non-copyable, non-movable
    GILAcquireGuard(const GILAcquireGuard&) = delete;
    GILAcquireGuard& operator=(const GILAcquireGuard&) = delete;

private:
    PyGILState_STATE gstate_;
    bool acquired_;
};

}  // namespace

// ============================================================================
// CUDA user-object deferred cleanup
//
// CUDA invokes a user-object destructor on an internal thread where CUDA
// calls are forbidden. Payload cleanup can release resource handles whose
// deleters call CUDA, so the callback only transfers a preallocated intrusive
// node to this process-lifetime queue. One coalesced pending call drains all
// queued payloads from Python's main thread.
// ============================================================================

// Intrusive base for payloads transferred out of CUDA's callback.
struct DeferredCleanupItem {
    DeferredCleanupItem* next = nullptr;
    virtual ~DeferredCleanupItem() noexcept = default;
};

namespace {

// Process-lifetime MPSC queue that drains payloads from Python's main thread.
class DeferredCleanupQueue {
public:
    // Transfer one preallocated cleanup item from a producer to the queue.
    void enqueue(DeferredCleanupItem* item) noexcept {
        DeferredCleanupItem* head = head_.load(std::memory_order_relaxed);
        do {
            item->next = head;
        } while (!head_.compare_exchange_weak(
            head, item, std::memory_order_release, std::memory_order_relaxed));
        schedule();
    }

    // Permanently disable pending-call scheduling during interpreter shutdown.
    void stop() noexcept {
        accepting_.store(false, std::memory_order_release);
    }

    // Reattempt scheduling after Py_AddPendingCall() found CPython's bounded
    // pending-call queue full and left payloads queued for a later safe entry.
    void retry_schedule() noexcept {
        schedule();
    }

private:
    // Adapt queue draining to CPython's int (*)(void*) callback ABI.
    static int pending_call(void* arg) noexcept {
        static_cast<DeferredCleanupQueue*>(arg)->drain();
        return 0;
    }

    // Coalesce all queued work behind at most one CPython pending call.
    void schedule() noexcept {
        if (!accepting_.load(std::memory_order_acquire)) {
            return;
        }
        if (!Py_IsInitialized() || py_is_finalizing()) {
            stop();
            return;
        }
        if (!head_.load(std::memory_order_acquire)) {
            return;
        }
        bool expected = false;
        if (!scheduled_.compare_exchange_strong(
                expected, true, std::memory_order_acq_rel,
                std::memory_order_relaxed)) {
            return;
        }
        if (Py_AddPendingCall(&DeferredCleanupQueue::pending_call, this) != 0) {
            // Keep every payload queued. A later enqueue or safe cuda-core
            // entry can retry without blocking CUDA's callback thread.
            scheduled_.store(false, std::memory_order_release);
        }
    }

    // Detach and destroy all queued payloads from Python's main thread.
    void drain() noexcept {
        if (!Py_IsInitialized() || py_is_finalizing()) {
            stop();
            scheduled_.store(false, std::memory_order_release);
            return;  // Intentionally leak intact payloads during shutdown.
        }

        while (DeferredCleanupItem* list =
                   head_.exchange(nullptr, std::memory_order_acquire)) {
            while (list) {
                DeferredCleanupItem* next = list->next;
                delete list;
                list = next;
            }
        }

        scheduled_.store(false, std::memory_order_release);
        if (head_.load(std::memory_order_acquire)) {
            schedule();
        }
    }

    // Head of the intrusive multi-producer, single-consumer payload stack.
    std::atomic<DeferredCleanupItem*> head_{nullptr};
    // True while one cuda-core drain callback is pending or executing.
    std::atomic<bool> scheduled_{false};
    // False once shutdown begins, causing later payloads to be leaked safely.
    std::atomic<bool> accepting_{true};
};

// Published once at module initialization and intentionally never freed.
std::atomic<DeferredCleanupQueue*> deferred_cleanup_queue{nullptr};

void ensure_deferred_cleanup_ready() {
    DeferredCleanupQueue* queue =
        deferred_cleanup_queue.load(std::memory_order_acquire);
    if (!queue) {
        throw std::runtime_error("deferred cleanup is not initialized");
    }
    queue->retry_schedule();
}

// CUDA's CUhostFn ABI is void (*)(void*); recover and enqueue the cleanup item.
void enqueue_cleanup(void* item) noexcept {
    auto* cleanup = static_cast<DeferredCleanupItem*>(item);
    if (DeferredCleanupQueue* queue =
            deferred_cleanup_queue.load(std::memory_order_acquire)) {
        queue->enqueue(cleanup);
    }
}

}  // namespace

// Module initialization calls this once with the GIL held, which serializes
// the check, allocation, and publication below.
void initialize_deferred_cleanup() {
    if (deferred_cleanup_queue.load(std::memory_order_acquire)) {
        return;
    }
    auto* queue = new DeferredCleanupQueue();
    deferred_cleanup_queue.store(queue, std::memory_order_release);
}

void retry_deferred_cleanup() noexcept {
    if (!Py_IsInitialized() || py_is_finalizing()) {
        return;
    }
    if (DeferredCleanupQueue* queue =
            deferred_cleanup_queue.load(std::memory_order_acquire)) {
        queue->retry_schedule();
    }
}

// ============================================================================
// Handle reverse-lookup registry
//
// Maps raw CUDA handles (CUevent, CUkernel, etc.) back to their owning
// shared_ptr so that _ref constructors can recover full metadata.
// Uses weak_ptr to avoid preventing destruction.
// ============================================================================

template<typename Key, typename Handle, typename Hash = std::hash<Key>>
class HandleRegistry {
public:
    using MapType = std::unordered_map<Key, std::weak_ptr<typename Handle::element_type>, Hash>;

    void register_handle(const Key& key, const Handle& h) {
        std::lock_guard<std::mutex> lock(mutex_);
        map_[key] = h;
    }

    void unregister_handle(const Key& key) noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        map_.erase(key);
    }

    void register_handles(const std::vector<Handle>& handles) {
        std::lock_guard<std::mutex> lock(mutex_);
        for (const Handle& h : handles) {
            if (h) {
                map_[*h] = h;
            }
        }
    }

    Handle lookup(const Key& key) {
        std::lock_guard<std::mutex> lock(mutex_);
        auto it = map_.find(key);
        if (it != map_.end()) {
            if (auto h = it->second.lock()) {
                return h;
            }
            map_.erase(it);
        }
        return {};
    }

    template<typename Factory>
    Handle get_or_create(const Key& key, Factory&& create) {
        std::lock_guard<std::mutex> lock(mutex_);
        auto it = map_.find(key);
        if (it != map_.end()) {
            if (Handle h = it->second.lock()) {
                return h;
            }
            map_.erase(it);
        }

        Handle h = create();
        if (h) {
            map_[key] = h;
        }
        return h;
    }

    MapType drain() noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        MapType extracted;
        extracted.swap(map_);
        return extracted;
    }

private:
    std::mutex mutex_;
    MapType map_;
};

// ============================================================================
// Thread-local error handling
// ============================================================================

// Thread-local status of the most recent CUDA API call in this module.
static thread_local CUresult err = CUDA_SUCCESS;

CUresult get_last_error() noexcept {
    CUresult e = err;
    err = CUDA_SUCCESS;
    return e;
}

CUresult peek_last_error() noexcept {
    return err;
}

void clear_last_error() noexcept {
    err = CUDA_SUCCESS;
}

// ============================================================================
// Context Handles
// ============================================================================

namespace {
struct ContextBox {
    CUcontext resource;
    GreenCtxHandle h_green_ctx;
};

struct GreenCtxBox {
    CUgreenCtx resource;
};

static const ContextBox* get_box(const ContextHandle& h) noexcept {
    const CUcontext* p = h.get();
    return reinterpret_cast<const ContextBox*>(
        reinterpret_cast<const char*>(p) - offsetof(ContextBox, resource)
    );
}

// See REGISTRY_DESIGN.md (Level 1: Driver Handle -> Resource Handle)
static HandleRegistry<CUcontext, ContextHandle> context_registry;

// Create a context handle reference, with optional green context as source.
ContextHandle create_context_handle_ref(CUcontext ctx, GreenCtxHandle h_green_ctx) {
    if (!ctx) {
        return {};
    }
    if (auto h = context_registry.lookup(ctx)) {
        return h;
    }
    auto box = std::shared_ptr<const ContextBox>(
        new ContextBox{ctx, std::move(h_green_ctx)},
        [](const ContextBox* b) {
            context_registry.unregister_handle(b->resource);
            delete b;
        }
    );
    ContextHandle h(box, &box->resource);
    context_registry.register_handle(ctx, h);
    return h;
}
}  // namespace

ContextHandle create_context_handle_ref(CUcontext ctx) {
    return create_context_handle_ref(ctx, {});
}

ContextHandle create_context_handle_from_green_ctx(const GreenCtxHandle& h_green_ctx) {
    GILReleaseGuard gil;
    if (!h_green_ctx) {
        return {};
    }
    if (!p_cuCtxFromGreenCtx) {
        err = CUDA_ERROR_NOT_SUPPORTED;
        return {};
    }

    CUcontext ctx = nullptr;
    if (CUDA_SUCCESS != (err = p_cuCtxFromGreenCtx(&ctx, as_cu(h_green_ctx)))) {
        return {};
    }

    return create_context_handle_ref(ctx, h_green_ctx);
}

GreenCtxHandle get_context_green_ctx(const ContextHandle& h) noexcept {
    if (!h) {
        return {};
    }
    return get_box(h)->h_green_ctx;
}

GreenCtxHandle create_green_ctx_handle(CUdevResource* resources, unsigned int nbResources,
                                       CUdevice dev, unsigned int flags) {
    GILReleaseGuard gil;
    if (!p_cuDevResourceGenerateDesc || !p_cuGreenCtxCreate || !p_cuGreenCtxDestroy) {
        err = CUDA_ERROR_NOT_SUPPORTED;
        return {};
    }

    CUdevResourceDesc desc = nullptr;
    if (CUDA_SUCCESS != (err = p_cuDevResourceGenerateDesc(&desc, resources, nbResources))) {
        return {};
    }

    CUgreenCtx green_ctx = nullptr;
    if (CUDA_SUCCESS != (err = p_cuGreenCtxCreate(&green_ctx, desc, dev, flags))) {
        return {};
    }

    auto box = std::shared_ptr<const GreenCtxBox>(
        new GreenCtxBox{green_ctx},
        [](const GreenCtxBox* b) {
            GILReleaseGuard gil;
            p_cuGreenCtxDestroy(b->resource);
            delete b;
        }
    );
    return GreenCtxHandle(box, &box->resource);
}

GreenCtxHandle create_green_ctx_handle_ref(CUgreenCtx green_ctx) {
    if (!green_ctx) {
        return {};
    }
    auto box = std::make_shared<const GreenCtxBox>(GreenCtxBox{green_ctx});
    return GreenCtxHandle(box, &box->resource);
}

// Thread-local cache of primary contexts indexed by device ID
static thread_local std::vector<ContextHandle> primary_context_cache;

ContextHandle get_primary_context(int device_id) {
    // Check thread-local cache
    if (static_cast<size_t>(device_id) < primary_context_cache.size()) {
        if (auto cached = primary_context_cache[device_id]) {
            return cached;
        }
    }

    // Cache miss - acquire primary context from driver
    GILReleaseGuard gil;
    CUcontext ctx;
    if (CUDA_SUCCESS != (err = p_cuDevicePrimaryCtxRetain(&ctx, device_id))) {
        return {};
    }

    auto box = std::shared_ptr<const ContextBox>(
        new ContextBox{ctx, {}},
        [device_id](const ContextBox* b) {
            context_registry.unregister_handle(b->resource);
            GILReleaseGuard gil;
            p_cuDevicePrimaryCtxRelease(device_id);
            delete b;
        }
    );
    auto h = ContextHandle(box, &box->resource);
    context_registry.register_handle(ctx, h);

    // Update cache
    if (static_cast<size_t>(device_id) >= primary_context_cache.size()) {
        primary_context_cache.resize(device_id + 1);
    }
    primary_context_cache[device_id] = h;
    return h;
}

ContextHandle get_current_context() {
    GILReleaseGuard gil;
    CUcontext ctx = nullptr;
    if (CUDA_SUCCESS != (err = p_cuCtxGetCurrent(&ctx))) {
        return {};
    }
    if (!ctx) {
        return {};  // No current context (not an error)
    }
    return create_context_handle_ref(ctx);
}

// ============================================================================
// Stream Handles
// ============================================================================

namespace {
struct StreamBox {
    CUstream resource;
    ContextHandle h_context;
};

static const StreamBox* get_box(const StreamHandle& h) noexcept {
    const CUstream* p = h.get();
    return reinterpret_cast<const StreamBox*>(
        reinterpret_cast<const char*>(p) - offsetof(StreamBox, resource)
    );
}

// See REGISTRY_DESIGN.md (Level 1: Driver Handle -> Resource Handle)
static HandleRegistry<CUstream, StreamHandle> stream_registry;
}  // namespace

StreamHandle create_stream_handle(const ContextHandle& h_ctx, unsigned int flags, int priority) {
    GILReleaseGuard gil;
    CUstream stream;

    // Dispatch: green context uses cuGreenCtxStreamCreate, primary uses cuStreamCreateWithPriority
    GreenCtxHandle h_green = get_context_green_ctx(h_ctx);
    if (h_green) {
        if (!p_cuGreenCtxStreamCreate) {
            err = CUDA_ERROR_NOT_SUPPORTED;
            return {};
        }
        if (CUDA_SUCCESS != (err = p_cuGreenCtxStreamCreate(&stream, as_cu(h_green), flags, priority))) {
            return {};
        }
    } else {
        if (CUDA_SUCCESS != (err = p_cuStreamCreateWithPriority(&stream, flags, priority))) {
            return {};
        }
    }

    auto box = std::shared_ptr<const StreamBox>(
        new StreamBox{stream, h_ctx},
        [](const StreamBox* b) {
            stream_registry.unregister_handle(b->resource);
            GILReleaseGuard gil;
            p_cuStreamDestroy(b->resource);
            delete b;
        }
    );
    StreamHandle h(box, &box->resource);
    stream_registry.register_handle(stream, h);
    return h;
}

StreamHandle create_stream_handle_ref(CUstream stream) {
    if (auto h = stream_registry.lookup(stream)) {
        return h;
    }
    auto box = std::shared_ptr<const StreamBox>(
        new StreamBox{stream, {}},
        [](const StreamBox* b) {
            stream_registry.unregister_handle(b->resource);
            delete b;
        }
    );
    StreamHandle h(box, &box->resource);
    stream_registry.register_handle(stream, h);
    return h;
}

StreamHandle create_stream_handle_with_owner(CUstream stream, PyObject* owner) {
    if (auto h = stream_registry.lookup(stream)) {
        // Reuse handles that already carry structural context metadata, e.g.
        // cuda-core-owned streams.
        if (get_box(h)->h_context) {
            return h;
        }
    }
    if (!owner) {
        return create_stream_handle_ref(stream);
    }
    // GIL required when owner is provided
    GILAcquireGuard gil;
    if (!gil.acquired()) {
        // Python finalizing - fall back to ref version (no owner tracking)
        return create_stream_handle_ref(stream);
    }
    Py_INCREF(owner);
    // Owner-backed handles are NOT registered in the stream registry to avoid
    // corruption when multiple owners wrap the same CUstream (each stacks its
    // own Py_INCREF/Py_DECREF independently).
    auto box = std::shared_ptr<const StreamBox>(
        new StreamBox{stream, {}},
        [owner](const StreamBox* b) {
            GILAcquireGuard gil;
            if (gil.acquired()) {
                Py_DECREF(owner);
            }
            delete b;
        }
    );
    return StreamHandle(box, &box->resource);
}

void py_object_user_object_destroy(void* py_object) noexcept {
    if (!py_object) {
        return;
    }
    GILAcquireGuard gil;
    if (!gil.acquired()) {
        return;
    }
    Py_DECREF(reinterpret_cast<PyObject*>(py_object));
}

ContextHandle get_stream_context(const StreamHandle& h) noexcept {
    return h ? get_box(h)->h_context : ContextHandle{};
}

StreamHandle get_legacy_stream() {
    static StreamHandle handle = create_stream_handle_ref(CU_STREAM_LEGACY);
    return handle;
}

StreamHandle get_per_thread_stream() {
    static StreamHandle handle = create_stream_handle_ref(CU_STREAM_PER_THREAD);
    return handle;
}

// ============================================================================
// Event Handles
// ============================================================================

namespace {
struct EventBox {
    CUevent resource;
    bool timing_enabled;
    bool is_blocking_sync;
    bool ipc_enabled;
    int device_id;
    ContextHandle h_context;
};
}  // namespace

static const EventBox* get_box(const EventHandle& h) {
    const CUevent* p = h.get();
    return reinterpret_cast<const EventBox*>(
        reinterpret_cast<const char*>(p) - offsetof(EventBox, resource)
    );
}

bool get_event_timing_enabled(const EventHandle& h) noexcept {
    return h ? get_box(h)->timing_enabled : false;
}

bool get_event_is_blocking_sync(const EventHandle& h) noexcept {
    return h ? get_box(h)->is_blocking_sync : false;
}

bool get_event_ipc_enabled(const EventHandle& h) noexcept {
    return h ? get_box(h)->ipc_enabled : false;
}

int get_event_device_id(const EventHandle& h) noexcept {
    return h ? get_box(h)->device_id : -1;
}

ContextHandle get_event_context(const EventHandle& h) noexcept {
    return h ? get_box(h)->h_context : ContextHandle{};
}

// See REGISTRY_DESIGN.md (Level 1: Driver Handle -> Resource Handle)
static HandleRegistry<CUevent, EventHandle> event_registry;

EventHandle create_event_handle(const ContextHandle& h_ctx, unsigned int flags,
                                bool timing_enabled, bool is_blocking_sync,
                                bool ipc_enabled, int device_id) {
    GILReleaseGuard gil;
    CUevent event;
    if (CUDA_SUCCESS != (err = p_cuEventCreate(&event, flags))) {
        return {};
    }

    auto box = std::shared_ptr<const EventBox>(
        new EventBox{event, timing_enabled, is_blocking_sync, ipc_enabled, device_id, h_ctx},
        [h_ctx](const EventBox* b) {
            event_registry.unregister_handle(b->resource);
            GILReleaseGuard gil;
            p_cuEventDestroy(b->resource);
            delete b;
        }
    );
    EventHandle h(box, &box->resource);
    event_registry.register_handle(event, h);
    return h;
}

EventHandle create_event_handle_noctx(unsigned int flags) {
    return create_event_handle(ContextHandle{}, flags, false, false, false, -1);
}

EventHandle create_event_handle_ref(CUevent event) {
    if (auto h = event_registry.lookup(event)) {
        return h;
    }
    auto box = std::make_shared<const EventBox>(EventBox{event, false, false, false, -1, {}});
    return EventHandle(box, &box->resource);
}

EventHandle create_event_handle_ipc(const CUipcEventHandle& ipc_handle,
                                    bool is_blocking_sync) {
    GILReleaseGuard gil;
    CUevent event;
    if (CUDA_SUCCESS != (err = p_cuIpcOpenEventHandle(&event, ipc_handle))) {
        return {};
    }

    auto box = std::shared_ptr<const EventBox>(
        new EventBox{event, false, is_blocking_sync, true, -1, {}},
        [](const EventBox* b) {
            event_registry.unregister_handle(b->resource);
            GILReleaseGuard gil;
            p_cuEventDestroy(b->resource);
            delete b;
        }
    );
    EventHandle h(box, &box->resource);
    event_registry.register_handle(event, h);
    return h;
}

// ============================================================================
// Memory Pool Handles
// ============================================================================

namespace {
struct MemoryPoolBox {
    CUmemoryPool resource;
};
}  // namespace

// Helper to clear peer access before destroying a memory pool.
// Works around nvbug 5698116: recycled pool handles inherit peer access state.
// Must be noexcept since it's called from a shared_ptr deleter.
static void clear_mempool_peer_access(CUmemoryPool pool) noexcept {
    try {
        int device_count = 0;
        if (p_cuDeviceGetCount(&device_count) != CUDA_SUCCESS || device_count <= 0) {
            return;
        }

        std::vector<CUmemAccessDesc> clear_access(device_count);
        for (int i = 0; i < device_count; ++i) {
            clear_access[i].location.type = CU_MEM_LOCATION_TYPE_DEVICE;
            clear_access[i].location.id = i;
            clear_access[i].flags = CU_MEM_ACCESS_FLAGS_PROT_NONE;
        }
        p_cuMemPoolSetAccess(pool, clear_access.data(), device_count);  // Best effort
    } catch (...) {
        // Swallow exceptions - this is best-effort cleanup in destructor context
    }
}

static MemoryPoolHandle wrap_mempool_owned(CUmemoryPool pool) {
    auto box = std::shared_ptr<const MemoryPoolBox>(
        new MemoryPoolBox{pool},
        [](const MemoryPoolBox* b) {
            GILReleaseGuard gil;
            clear_mempool_peer_access(b->resource);
            p_cuMemPoolDestroy(b->resource);
            delete b;
        }
    );
    return MemoryPoolHandle(box, &box->resource);
}

MemoryPoolHandle create_mempool_handle(const CUmemPoolProps& props) {
    GILReleaseGuard gil;
    CUmemoryPool pool;
    if (CUDA_SUCCESS != (err = p_cuMemPoolCreate(&pool, &props))) {
        return {};
    }
    return wrap_mempool_owned(pool);
}

MemoryPoolHandle create_mempool_handle_ref(CUmemoryPool pool) {
    auto box = std::make_shared<const MemoryPoolBox>(MemoryPoolBox{pool});
    return MemoryPoolHandle(box, &box->resource);
}

MemoryPoolHandle get_device_mempool(int device_id) {
    GILReleaseGuard gil;
    CUmemoryPool pool;
    if (CUDA_SUCCESS != (err = p_cuDeviceGetMemPool(&pool, device_id))) {
        return {};
    }
    return create_mempool_handle_ref(pool);
}

MemoryPoolHandle create_mempool_handle_ipc(int fd, CUmemAllocationHandleType handle_type) {
    GILReleaseGuard gil;
    CUmemoryPool pool;
    auto handle_ptr = reinterpret_cast<void*>(static_cast<uintptr_t>(fd));
    if (CUDA_SUCCESS != (err = p_cuMemPoolImportFromShareableHandle(&pool, handle_ptr, handle_type, 0))) {
        return {};
    }
    return wrap_mempool_owned(pool);
}

// ============================================================================
// Device Pointer Handles
// ============================================================================

namespace {
struct DevicePtrBox {
    CUdeviceptr resource;
    // Mutable to allow set_deallocation_stream() to update the stream
    // through a const DevicePtrHandle. The stream can be changed after
    // allocation (e.g., to synchronize deallocation with a different stream).
    mutable StreamHandle h_stream;
};
}  // namespace

// Recovers the owning DevicePtrBox from the aliased CUdeviceptr pointer.
// This works because DevicePtrHandle is a shared_ptr alias pointing to
// &box->resource, so we can compute the containing struct using offsetof.
// The const_cast is safe because we only use this to access the mutable
// h_stream member or in the deleter (where the box is being destroyed).
static DevicePtrBox* get_box(const DevicePtrHandle& h) {
    const CUdeviceptr* p = h.get();
    return reinterpret_cast<DevicePtrBox*>(
        reinterpret_cast<char*>(const_cast<CUdeviceptr*>(p)) - offsetof(DevicePtrBox, resource)
    );
}

StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept {
    return get_box(h)->h_stream;
}

void set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept {
    get_box(h)->h_stream = h_stream;
}

DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h_pool, const StreamHandle& h_stream) {
    GILReleaseGuard gil;
    CUdeviceptr ptr;
    if (CUDA_SUCCESS != (err = p_cuMemAllocFromPoolAsync(&ptr, size, *h_pool, as_cu(h_stream)))) {
        return {};
    }

    auto box = std::shared_ptr<DevicePtrBox>(
        new DevicePtrBox{ptr, h_stream},
        [h_pool](DevicePtrBox* b) {
            GILReleaseGuard gil;
            p_cuMemFreeAsync(b->resource, as_cu(b->h_stream));
            delete b;
        }
    );
    return DevicePtrHandle(box, &box->resource);
}

DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream) {
    GILReleaseGuard gil;
    CUdeviceptr ptr;
    if (CUDA_SUCCESS != (err = p_cuMemAllocAsync(&ptr, size, as_cu(h_stream)))) {
        return {};
    }

    auto box = std::shared_ptr<DevicePtrBox>(
        new DevicePtrBox{ptr, h_stream},
        [](DevicePtrBox* b) {
            GILReleaseGuard gil;
            p_cuMemFreeAsync(b->resource, as_cu(b->h_stream));
            delete b;
        }
    );
    return DevicePtrHandle(box, &box->resource);
}

DevicePtrHandle deviceptr_alloc(size_t size) {
    GILReleaseGuard gil;
    CUdeviceptr ptr;
    if (CUDA_SUCCESS != (err = p_cuMemAlloc(&ptr, size))) {
        return {};
    }

    auto box = std::shared_ptr<DevicePtrBox>(
        new DevicePtrBox{ptr, StreamHandle{}},
        [](DevicePtrBox* b) {
            GILReleaseGuard gil;
            p_cuMemFree(b->resource);
            delete b;
        }
    );
    return DevicePtrHandle(box, &box->resource);
}

DevicePtrHandle deviceptr_alloc_host(size_t size) {
    GILReleaseGuard gil;
    void* ptr;
    if (CUDA_SUCCESS != (err = p_cuMemAllocHost(&ptr, size))) {
        return {};
    }

    auto box = std::shared_ptr<DevicePtrBox>(
        new DevicePtrBox{reinterpret_cast<CUdeviceptr>(ptr), StreamHandle{}},
        [](DevicePtrBox* b) {
            GILReleaseGuard gil;
            p_cuMemFreeHost(reinterpret_cast<void*>(b->resource));
            delete b;
        }
    );
    return DevicePtrHandle(box, &box->resource);
}

DevicePtrHandle deviceptr_create_ref(CUdeviceptr ptr) {
    auto box = std::make_shared<DevicePtrBox>(DevicePtrBox{ptr, StreamHandle{}});
    return DevicePtrHandle(box, &box->resource);
}

DevicePtrHandle deviceptr_create_with_owner(CUdeviceptr ptr, PyObject* owner) {
    if (!owner) {
        return deviceptr_create_ref(ptr);
    }
    // GIL required when owner is provided
    GILAcquireGuard gil;
    if (!gil.acquired()) {
        // Python finalizing - fall back to ref version (no owner tracking)
        return deviceptr_create_ref(ptr);
    }
    Py_INCREF(owner);
    auto box = std::shared_ptr<DevicePtrBox>(
        new DevicePtrBox{ptr, StreamHandle{}},
        [owner](DevicePtrBox* b) {
            GILAcquireGuard gil;
            if (gil.acquired()) {
                Py_DECREF(owner);
            }
            delete b;
        }
    );
    return DevicePtrHandle(box, &box->resource);
}

DevicePtrHandle deviceptr_create_mapped_graphics(
    CUdeviceptr ptr,
    const GraphicsResourceHandle& h_resource,
    const StreamHandle& h_stream
) {
    auto box = std::shared_ptr<DevicePtrBox>(
        new DevicePtrBox{ptr, h_stream},
        [h_resource](DevicePtrBox* b) {
            GILReleaseGuard gil;
            CUgraphicsResource resource = as_cu(h_resource);
            p_cuGraphicsUnmapResources(1, &resource, as_cu(b->h_stream));
            delete b;
        }
    );
    return DevicePtrHandle(box, &box->resource);
}

// ============================================================================
// MemoryResource-owned Device Pointer Handles
// ============================================================================

static MRDeallocCallback mr_dealloc_cb = nullptr;

void register_mr_dealloc_callback(MRDeallocCallback cb) {
    mr_dealloc_cb = cb;
}

DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* mr) {
    if (!mr) {
        return deviceptr_create_ref(ptr);
    }
    // GIL required when mr is provided
    GILAcquireGuard gil;
    if (!gil.acquired()) {
        return deviceptr_create_ref(ptr);
    }
    Py_INCREF(mr);
    auto box = std::shared_ptr<DevicePtrBox>(
        new DevicePtrBox{ptr, StreamHandle{}},
        [mr, size](DevicePtrBox* b) {
            GILAcquireGuard gil;
            if (gil.acquired()) {
                if (mr_dealloc_cb) {
                    mr_dealloc_cb(mr, b->resource, size, b->h_stream);
                }
                Py_DECREF(mr);
            }
            delete b;
        }
    );
    return DevicePtrHandle(box, &box->resource);
}

// ============================================================================
// IPC Pointer Cache
// ============================================================================
// This cache handles duplicate IPC imports, which behave differently depending
// on the memory type:
//
// 1. Memory pool allocations (DeviceMemoryResource):
//    Multiple imports of the same allocation succeed and return duplicate
//    pointers. However, the driver has a reference counting bug (nvbug 5570902)
//    where the first cuMemFreeAsync incorrectly unmaps the memory even when
//    imported multiple times. A driver fix is expected.
//
// 2. Pinned memory allocations (PinnedMemoryResource):
//    Duplicate imports result in CUDA_ERROR_ALREADY_MAPPED.
//
// The cache solves both issues by checking the cache before calling
// cuMemPoolImportPointer and returning the existing handle for duplicate
// imports. This provides a consistent user experience where the same IPC
// descriptor can be imported multiple times regardless of memory type.
//
// The cache key is the export_data bytes (CUmemPoolPtrExportData), not the
// returned pointer, because we must check before calling the driver API.


// TODO: When driver fix for nvbug 5570902 is available, consider whether
// the cache is still needed for memory pool allocations (it will still be
// needed for pinned memory).
static bool use_ipc_ptr_cache() {
    return true;
}

namespace {
// Wrapper for CUmemPoolPtrExportData to use as map key
struct ExportDataKey {
    CUmemPoolPtrExportData data;

    bool operator==(const ExportDataKey& other) const {
        return std::memcmp(&data, &other.data, sizeof(data)) == 0;
    }
};

struct ExportDataKeyHash {
    std::size_t operator()(const ExportDataKey& key) const {
        // Simple hash of the bytes
        std::size_t h = 0;
        const auto* bytes = reinterpret_cast<const unsigned char*>(&key.data);
        for (std::size_t i = 0; i < sizeof(key.data); ++i) {
            h = h * 31 + bytes[i];
        }
        return h;
    }
};

}

static HandleRegistry<ExportDataKey, DevicePtrHandle, ExportDataKeyHash> ipc_ptr_cache;
static std::mutex ipc_import_mutex;

DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* export_data, const StreamHandle& h_stream) {
    auto data = const_cast<CUmemPoolPtrExportData*>(
        reinterpret_cast<const CUmemPoolPtrExportData*>(export_data));

    if (use_ipc_ptr_cache()) {
        ExportDataKey key;
        std::memcpy(&key.data, data, sizeof(key.data));

        std::lock_guard<std::mutex> lock(ipc_import_mutex);

        if (auto h = ipc_ptr_cache.lookup(key)) {
            return h;
        }

        GILReleaseGuard gil;
        CUdeviceptr ptr;
        if (CUDA_SUCCESS != (err = p_cuMemPoolImportPointer(&ptr, *h_pool, data))) {
            return {};
        }

        auto box = std::shared_ptr<DevicePtrBox>(
            new DevicePtrBox{ptr, h_stream},
            [h_pool, key](DevicePtrBox* b) {
                ipc_ptr_cache.unregister_handle(key);
                GILReleaseGuard gil;
                p_cuMemFreeAsync(b->resource, as_cu(b->h_stream));
                delete b;
            }
        );
        DevicePtrHandle h(box, &box->resource);
        ipc_ptr_cache.register_handle(key, h);
        return h;

    } else {
        GILReleaseGuard gil;
        CUdeviceptr ptr;
        if (CUDA_SUCCESS != (err = p_cuMemPoolImportPointer(&ptr, *h_pool, data))) {
            return {};
        }

        auto box = std::shared_ptr<DevicePtrBox>(
            new DevicePtrBox{ptr, h_stream},
            [h_pool](DevicePtrBox* b) {
                GILReleaseGuard gil;
                p_cuMemFreeAsync(b->resource, as_cu(b->h_stream));
                delete b;
            }
        );
        return DevicePtrHandle(box, &box->resource);
    }
}

// ============================================================================
// Library Handles
// ============================================================================

namespace {
struct LibraryBox {
    CUlibrary resource;
};
}  // namespace

LibraryHandle create_library_handle_from_file(const char* path) {
    GILReleaseGuard gil;
    CUlibrary library;
    if (CUDA_SUCCESS != (err = p_cuLibraryLoadFromFile(&library, path, nullptr, nullptr, 0, nullptr, nullptr, 0))) {
        return {};
    }

    auto box = std::shared_ptr<const LibraryBox>(
        new LibraryBox{library},
        [](const LibraryBox* b) {
            GILReleaseGuard gil;
            // TODO: re-enable once LibraryBox tracks its owning context
            // p_cuLibraryUnload(b->resource);
            delete b;
        }
    );
    return LibraryHandle(box, &box->resource);
}

LibraryHandle create_library_handle_from_data(const void* data) {
    GILReleaseGuard gil;
    CUlibrary library;
    if (CUDA_SUCCESS != (err = p_cuLibraryLoadData(&library, data, nullptr, nullptr, 0, nullptr, nullptr, 0))) {
        return {};
    }

    auto box = std::shared_ptr<const LibraryBox>(
        new LibraryBox{library},
        [](const LibraryBox* b) {
            GILReleaseGuard gil;
            // TODO: re-enable once LibraryBox tracks its owning context
            // p_cuLibraryUnload(b->resource);
            delete b;
        }
    );
    return LibraryHandle(box, &box->resource);
}

LibraryHandle create_library_handle_ref(CUlibrary library) {
    auto box = std::make_shared<const LibraryBox>(LibraryBox{library});
    return LibraryHandle(box, &box->resource);
}

// ============================================================================
// Kernel Handles
// ============================================================================

namespace {
struct KernelBox {
    CUkernel resource;
    LibraryHandle h_library;
};
}  // namespace

static const KernelBox* get_box(const KernelHandle& h) {
    const CUkernel* p = h.get();
    return reinterpret_cast<const KernelBox*>(
        reinterpret_cast<const char*>(p) - offsetof(KernelBox, resource)
    );
}

// See REGISTRY_DESIGN.md (Level 1: Driver Handle -> Resource Handle)
static HandleRegistry<CUkernel, KernelHandle> kernel_registry;

KernelHandle create_kernel_handle(const LibraryHandle& h_library, const char* name) {
    GILReleaseGuard gil;
    CUkernel kernel;
    if (CUDA_SUCCESS != (err = p_cuLibraryGetKernel(&kernel, *h_library, name))) {
        return {};
    }

    auto box = std::make_shared<const KernelBox>(KernelBox{kernel, h_library});
    KernelHandle h(box, &box->resource);
    kernel_registry.register_handle(kernel, h);
    return h;
}

KernelHandle create_kernel_handle_ref(CUkernel kernel) {
    if (auto h = kernel_registry.lookup(kernel)) {
        return h;
    }
    auto box = std::make_shared<const KernelBox>(KernelBox{kernel, {}});
    return KernelHandle(box, &box->resource);
}

LibraryHandle get_kernel_library(const KernelHandle& h) noexcept {
    if (!h) return {};
    return get_box(h)->h_library;
}

// ============================================================================
// Graph Handles
// ============================================================================

namespace {

struct NodeAttachment;
using GraphAttachmentMap = std::map<CUgraphNode, NodeAttachment*>;

struct GraphHierarchy;

// Standard-layout alias target for GraphHandle.
struct GraphBoxBase {
    CUgraph resource = nullptr;
};

// Canonical state for one CUgraph. Its GraphHandle aliases resource, whose
// address remains stable for the lifetime of the hierarchy.
struct GraphBox : GraphBoxBase {
    GraphHierarchy* hierarchy = nullptr;  // Non-owning back-reference.
    GraphBox* parent = nullptr;           // Null for the root graph.
    CUgraphNode owner_node = nullptr;     // Node in parent that owns this graph.
    GraphAttachmentMap attachments;       // Non-owning attachment index.
    HandleRegistry<CUgraphNode, GraphNodeHandle> node_handles;

    GraphBox(
            CUgraph resource_,
            GraphHierarchy* hierarchy_,
            GraphBox* parent_ = nullptr,
            CUgraphNode owner_node_ = nullptr) noexcept
        : GraphBoxBase{resource_},
          hierarchy(hierarchy_),
          parent(parent_),
          owner_node(owner_node_) {}
};

// Shared owner of stable GraphBox storage. Every GraphHandle aliases the same
// control block, so any graph handle keeps the entire hierarchy alive.
struct GraphHierarchy {
    std::list<GraphBox> graphs;  // Parent boxes precede their descendants.
    std::list<GraphBox> graveyard;  // Retired child graph tombstones.

    GraphBox* root() noexcept {
        return graphs.empty() ? nullptr : &graphs.front();
    }
};

// See REGISTRY_DESIGN.md (Level 1: Driver Handle -> Resource Handle)
using GraphRegistry = HandleRegistry<CUgraph, GraphHandle>;
static GraphRegistry graph_registry;

// Immutable resource owners for one version of a graph node's parameters.
// Inheriting DeferredCleanupItem lets CUDA's user-object destructor enqueue
// the payload without destroying owners on the callback thread.
struct NodeAttachment : DeferredCleanupItem {
    CUuserObject object = nullptr;
    std::array<OpaqueHandle, 2> owners;

    NodeAttachment(OpaqueHandle owner0, OpaqueHandle owner1)
        : owners{std::move(owner0), std::move(owner1)} {}
};

// shared_ptr deleters for the payloads that need one. Typed handles convert to
// OpaqueHandle by assignment and reuse their own control block, so they need no
// deleter here. The Python deleter follows the owner-release pattern used by
// the stream/deviceptr handles above.
void py_deleter(const void* p) noexcept {
    GILAcquireGuard gil;
    if (gil.acquired()) {
        Py_DECREF(const_cast<PyObject*>(static_cast<const PyObject*>(p)));
    }
}

void free_deleter(const void* p) noexcept {
    std::free(const_cast<void*>(p));
}

GraphBox* get_box(const GraphHandle& h) noexcept {
    auto* value = reinterpret_cast<const GraphBoxBase*>(h.get());
    return const_cast<GraphBox*>(
        static_cast<const GraphBox*>(value));
}

// Rekey a staged attachment map from source nodes to their cloned nodes.
// The caller must release the GIL before calling this function.
CUresult rekey_attachments(
        GraphAttachmentMap& attachments, CUgraph cloned_graph) {
    if (!cloned_graph) {
        return CUDA_ERROR_INVALID_VALUE;
    }
    if (!p_cuGraphNodeFindInClone) {
        return CUDA_ERROR_NOT_SUPPORTED;
    }

    GraphAttachmentMap remapped;
    while (!attachments.empty()) {
        auto attachment = attachments.extract(attachments.begin());
        CUgraphNode cloned_node = nullptr;
        CUresult status = p_cuGraphNodeFindInClone(
            &cloned_node, attachment.key(), cloned_graph);
        if (status != CUDA_SUCCESS) {
            return status;
        }
        attachment.key() = cloned_node;
        if (!remapped.insert(std::move(attachment)).inserted) {
            return CUDA_ERROR_INVALID_VALUE;
        }
    }
    attachments.swap(remapped);
    return CUDA_SUCCESS;
}

struct StagedGraphMetadata {
    const GraphBox* source;
    GraphBox* clone;
    GraphAttachmentMap* attachments;
};
using StagedGraphMetadataList = std::vector<StagedGraphMetadata>;

// Copy a source hierarchy into detached metadata before CUDA mutation.
void stage_graph_metadata(
        const GraphBox& source,
        GraphBox& clone,
        GraphAttachmentMap& attachments,
        std::list<GraphBox>& subgraphs,
        StagedGraphMetadataList& staged) {
    attachments = source.attachments;
    staged.push_back({&source, &clone, &attachments});

    for (const GraphBox& source_child : source.hierarchy->graphs) {
        if (source_child.parent != &source || !source_child.resource) {
            continue;
        }
        GraphBox& cloned_child = subgraphs.emplace_back(
            nullptr,
            clone.hierarchy,
            &clone,
            nullptr);
        stage_graph_metadata(
            source_child,
            cloned_child,
            cloned_child.attachments,
            subgraphs,
            staged);
    }
}

// Bind staged metadata to a CUDA-cloned hierarchy. The root clone resource
// must be populated before entry. The caller must release the GIL.
CUresult rekey_graph_metadata(
        StagedGraphMetadataList& staged) {
    if (!p_cuGraphNodeFindInClone || !p_cuGraphChildGraphNodeGetGraph) {
        return CUDA_ERROR_NOT_SUPPORTED;
    }

    CUresult status;
    for (size_t i = 0; i < staged.size(); ++i) {
        const GraphBox& source = *staged[i].source;
        GraphBox& clone = *staged[i].clone;
        if (i != 0) {
            CUgraphNode cloned_owner = nullptr;
            status = p_cuGraphNodeFindInClone(
                &cloned_owner,
                source.owner_node,
                clone.parent->resource);
            if (status == CUDA_SUCCESS) {
                status = p_cuGraphChildGraphNodeGetGraph(
                    cloned_owner, &clone.resource);
            }
            if (status != CUDA_SUCCESS) {
                return status;
            }
            clone.owner_node = cloned_owner;
        }

        status = rekey_attachments(
            *staged[i].attachments, clone.resource);
        if (status != CUDA_SUCCESS) {
            return status;
        }
    }
    return CUDA_SUCCESS;
}

}  // namespace

OpaqueHandle make_opaque_py(PyObject* obj) {
    Py_INCREF(obj);
    return OpaqueHandle(static_cast<const void*>(obj), py_deleter);
}

OpaqueHandle make_opaque_malloc(void* buf) {
    return OpaqueHandle(static_cast<const void*>(buf), free_deleter);
}

// State held by PreparedAttachment between preparation and commit. It keeps the
// graph alive, tracks the graph-retained replacement, and holds a preallocated
// map entry so commit cannot allocate. Destroying PreparedAttachment rolls back
// the staged user-object retain unless graph_commit_attachment publishes it.
struct PreparedAttachmentState {
    GraphHandle h_graph;
    NodeAttachment* replacement = nullptr;
    GraphAttachmentMap::node_type replacement_entry;

    explicit PreparedAttachmentState(GraphHandle h_graph_)
        : h_graph(std::move(h_graph_)) {}
};

void rollback_prepared_attachment(
        PreparedAttachmentState* state) noexcept {
    if (!state) {
        return;
    }
    if (state->replacement) {
        GraphBox* box = get_box(state->h_graph);
        if (box->resource) {
            GILReleaseGuard gil;
            p_cuGraphReleaseUserObject(
                box->resource, state->replacement->object, 1);
        }
    }
    delete state;
}

// Detached metadata for a replacement embedded graph hierarchy. Preparation
// copies every attachment map and allocates every GraphBox before CUDA destroys
// the old embedded graph. Commit only rekeys and publishes it.
struct PreparedChildGraphUpdateState {
    GraphHandle h_parent;
    GraphHandle h_source;
    GraphBox* old_root = nullptr;
    CUgraphNode owner_node = nullptr;
    std::list<GraphBox> replacement;
    StagedGraphMetadataList staged;
    std::vector<GraphHandle> handles;

    PreparedChildGraphUpdateState(
            GraphHandle h_parent_,
            GraphHandle h_source_,
            GraphBox* old_root_,
            CUgraphNode owner_node_)
        : h_parent(std::move(h_parent_)),
          h_source(std::move(h_source_)),
          old_root(old_root_),
          owner_node(owner_node_) {}
};

GraphHandle create_graph_handle(CUgraph graph) {
    if (!graph) {
        return {};
    }

    auto hierarchy = std::shared_ptr<GraphHierarchy>(
        new GraphHierarchy{},
        [](GraphHierarchy* hierarchy) {
            for (const GraphBox& box : hierarchy->graphs) {
                if (box.resource) {
                    graph_registry.unregister_handle(box.resource);
                }
            }
            GraphBox* root = hierarchy->root();
            if (root && root->resource) {
                GILReleaseGuard gil;
                p_cuGraphDestroy(root->resource);
            }
            retry_deferred_cleanup();
            delete hierarchy;
        }
    );
    GraphBox& root = hierarchy->graphs.emplace_back(
        graph, hierarchy.get());

    GraphHandle h_graph(hierarchy, &root.resource);
    graph_registry.register_handle(graph, h_graph);
    return h_graph;
}

GraphHandle create_child_graph_handle(
        CUgraph child_graph, const GraphHandle& h_parent,
        CUgraphNode owner_node) {
    if (!child_graph || !h_parent || !owner_node) {
        return {};
    }
    if (GraphHandle h_graph = graph_registry.lookup(child_graph)) {
        return h_graph;
    }

    GraphBox* parent = get_box(h_parent);
    GraphHierarchy* hierarchy = parent->hierarchy;
    GraphBox& child = hierarchy->graphs.emplace_back(
        child_graph, hierarchy, parent, owner_node);

    GraphHandle h_child(h_parent, &child.resource);
    graph_registry.register_handle(child_graph, h_child);
    return h_child;
}

CUresult graph_prepare_child_graph_update(
        const GraphHandle& h_parent,
        const GraphHandle& h_old_child,
        CUgraphNode owner_node,
        const GraphHandle& h_source,
        PreparedChildGraphUpdate* out_prepared) {
    if (!h_parent || !h_old_child || !owner_node ||
        !h_source || !out_prepared) {
        return CUDA_ERROR_INVALID_VALUE;
    }
    out_prepared->reset();

    GraphBox* parent = get_box(h_parent);
    GraphBox* old_root = get_box(h_old_child);
    GraphBox* source = get_box(h_source);
    // A source from the destination hierarchy can include the old embedded
    // subtree whose raw node keys CUDA destroys during replacement.
    if (!parent->resource || !old_root->resource || !source->resource ||
        old_root->parent != parent ||
        old_root->owner_node != owner_node ||
        source->hierarchy == parent->hierarchy) {
        return CUDA_ERROR_INVALID_VALUE;
    }

    PreparedChildGraphUpdate prepared =
        std::make_shared<PreparedChildGraphUpdateState>(
            h_parent, h_source, old_root, owner_node);

    GraphBox& replacement_root =
        prepared->replacement.emplace_back(
            nullptr, parent->hierarchy, parent, owner_node);
    stage_graph_metadata(
        *source,
        replacement_root,
        replacement_root.attachments,
        prepared->replacement,
        prepared->staged);

    const size_t graph_count = prepared->staged.size();
    prepared->handles.reserve(graph_count);
    for (const StagedGraphMetadata& graph : prepared->staged) {
        prepared->handles.emplace_back(
            h_parent, &graph.clone->resource);
    }

    *out_prepared = std::move(prepared);
    return CUDA_SUCCESS;
}

void publish_child_graph_update(
        PreparedChildGraphUpdateState& state,
        GraphHandle* out_child) {
    GraphBox* parent = get_box(state.h_parent);
    parent->hierarchy->graphs.splice(
        parent->hierarchy->graphs.end(), state.replacement);
    *out_child = state.handles.front();
    graph_registry.register_handles(state.handles);
}

CUresult graph_commit_child_graph_update(
        PreparedChildGraphUpdate& prepared,
        GraphHandle* out_child) {
    if (!prepared || !out_child) {
        return CUDA_ERROR_INVALID_VALUE;
    }
    out_child->reset();

    PreparedChildGraphUpdateState& state = *prepared;
    GraphBox* parent = get_box(state.h_parent);
    if (!parent->resource || !state.old_root->resource) {
        prepared.reset();
        return CUDA_ERROR_INVALID_VALUE;
    }

    CUresult status = CUDA_ERROR_NOT_SUPPORTED;
    CUgraph cloned_root = nullptr;
    if (p_cuGraphChildGraphNodeGetGraph) {
        GILReleaseGuard gil;
        status = p_cuGraphChildGraphNodeGetGraph(
            state.owner_node, &cloned_root);
        if (status == CUDA_SUCCESS) {
            state.staged.front().clone->resource = cloned_root;
            status = rekey_graph_metadata(state.staged);
        }
    }

    // CUDA has already destroyed the old embedded graph. No replacement
    // metadata is visible yet, so this selects only the old generation.
    invalidate_child_graph_state(
        state.h_parent, state.owner_node);

    if (status != CUDA_SUCCESS) {
        prepared.reset();
        throw std::runtime_error(
            "failed to update graph metadata after child graph replacement");
    }

    publish_child_graph_update(state, out_child);
    prepared.reset();
    return status;
}

CUresult graph_get_attachment(
        const GraphHandle& h_graph, CUgraphNode node,
        OpaqueHandle* owner0, OpaqueHandle* owner1) {
    if (!h_graph || !node || (!owner0 && !owner1)) {
        return CUDA_ERROR_INVALID_VALUE;
    }
    if (owner0) {
        owner0->reset();
    }
    if (owner1) {
        owner1->reset();
    }

    GraphBox* box = get_box(h_graph);
    if (!box->resource) {
        return CUDA_ERROR_INVALID_VALUE;
    }
    auto it = box->attachments.find(node);
    if (it != box->attachments.end()) {
        if (owner0) {
            *owner0 = it->second->owners[0];
        }
        if (owner1) {
            *owner1 = it->second->owners[1];
        }
    }
    return CUDA_SUCCESS;
}

CUresult graph_prepare_attachment(
        const GraphHandle& h_graph,
        OpaqueHandle owner0,
        OpaqueHandle owner1,
        PreparedAttachment* out_prepared) {
    if (!out_prepared) {
        return CUDA_ERROR_INVALID_VALUE;
    }
    out_prepared->reset();
    if (!h_graph) {
        return CUDA_ERROR_INVALID_VALUE;
    }

    GraphBox* box = get_box(h_graph);
    if (!box->resource) {
        return CUDA_ERROR_INVALID_VALUE;
    }
    if (!p_cuGraphReleaseUserObject) {
        return CUDA_ERROR_NOT_SUPPORTED;
    }

    PreparedAttachment prepared(
        new PreparedAttachmentState(h_graph),
        PreparedAttachmentDeleter{rollback_prepared_attachment});
    if (owner0 || owner1) {
        if (!p_cuUserObjectCreate || !p_cuUserObjectRelease ||
            !p_cuGraphRetainUserObject) {
            return CUDA_ERROR_NOT_SUPPORTED;
        }

        ensure_deferred_cleanup_ready();
        prepared->replacement = new NodeAttachment(
            std::move(owner0), std::move(owner1));
        GraphAttachmentMap staged;
        try {
            staged.emplace(nullptr, prepared->replacement);
            prepared->replacement_entry =
                staged.extract(staged.begin());
        } catch (...) {
            delete prepared->replacement;
            prepared->replacement = nullptr;
            throw;
        }
        auto* cleanup_item =
            static_cast<DeferredCleanupItem*>(
                prepared->replacement);

        CUuserObject object = nullptr;
        CUresult status;
        {
            GILReleaseGuard gil;
            status = p_cuUserObjectCreate(
                &object, cleanup_item,
                reinterpret_cast<CUhostFn>(enqueue_cleanup),
                1, CU_USER_OBJECT_NO_DESTRUCTOR_SYNC);
            if (status != CUDA_SUCCESS) {
                prepared->replacement_entry.mapped() = nullptr;
                delete prepared->replacement;
                prepared->replacement = nullptr;
                return status;
            }
            prepared->replacement->object = object;
            status = p_cuGraphRetainUserObject(
                box->resource, object, 1, CU_GRAPH_USER_OBJECT_MOVE);
            if (status != CUDA_SUCCESS) {
                prepared->replacement_entry.mapped() = nullptr;
                prepared->replacement = nullptr;
                p_cuUserObjectRelease(object, 1);
                return status;
            }
        }
    }

    *out_prepared = std::move(prepared);
    return CUDA_SUCCESS;
}

CUresult graph_commit_attachment(
        PreparedAttachment& prepared,
        CUgraphNode node) {
    if (!prepared) {
        return CUDA_ERROR_INVALID_VALUE;
    }

    GraphHandle h_graph = prepared->h_graph;
    GraphBox* box = get_box(h_graph);
    if (!box->resource || (!node && !prepared->replacement)) {
        delete prepared.release();
        return CUDA_ERROR_INVALID_VALUE;
    }
    if (!node) {
        delete prepared.release();
        return CUDA_SUCCESS;
    }

    // Publish the replacement or removal before releasing the previous graph
    // reference; that release can make the previous payload eligible for
    // destruction.
    NodeAttachment* previous = nullptr;
    auto it = box->attachments.find(node);
    if (it == box->attachments.end()) {
        if (prepared->replacement) {
            prepared->replacement_entry.key() = node;
            auto result = box->attachments.insert(
                std::move(prepared->replacement_entry));
            if (!result.inserted) {
                prepared->replacement_entry =
                    std::move(result.node);
                delete prepared.release();
                return CUDA_ERROR_INVALID_VALUE;
            }
        }
    } else {
        previous = it->second;
        if (prepared->replacement) {
            it->second = prepared->replacement;
        } else {
            box->attachments.erase(it);
        }
    }

    delete prepared.release();
    if (!previous) {
        return CUDA_SUCCESS;
    }
    GILReleaseGuard gil;
    return p_cuGraphReleaseUserObject(
        box->resource, previous->object, 1);
}

CUresult graph_clone_attachments(
        const GraphHandle& h_clone,
        const GraphHandle& h_source) {
    if (!h_clone || !h_source) {
        return CUDA_ERROR_INVALID_VALUE;
    }

    GraphBox* clone = get_box(h_clone);
    GraphBox* source = get_box(h_source);
    if (!clone->resource || !source->resource ||
        !clone->attachments.empty()) {
        return CUDA_ERROR_INVALID_VALUE;
    }

    // Build and rekey the clone metadata off-hierarchy so a CUDA mapping error
    // cannot partially publish it.
    GraphAttachmentMap attachments;
    std::list<GraphBox> subgraphs;
    StagedGraphMetadataList staged;
    stage_graph_metadata(
        *source, *clone, attachments, subgraphs, staged);

    std::vector<GraphHandle> handles;
    handles.reserve(subgraphs.size());
    for (GraphBox& graph : subgraphs) {
        handles.emplace_back(h_clone, &graph.resource);
    }

    CUresult status;
    {
        GILReleaseGuard gil;
        status = rekey_graph_metadata(staged);
    }
    if (status != CUDA_SUCCESS) {
        return status;
    }

    clone->attachments.swap(attachments);
    if (subgraphs.empty()) {
        return CUDA_SUCCESS;
    }

    clone->hierarchy->graphs.splice(
        clone->hierarchy->graphs.end(), subgraphs);
    graph_registry.register_handles(handles);
    return CUDA_SUCCESS;
}

// ============================================================================
// Graph Exec Handles
// ============================================================================

namespace {

// Append-only owners introduced by individual executable-node updates. CUDA
// owns this payload through a user object propagated into the CUgraphExec.
struct ExecAttachments : DeferredCleanupItem {
    CUuserObject object = nullptr;
    std::vector<OpaqueHandle> owners;
};

struct GraphExecBox {
    CUgraphExec resource = nullptr;
    ExecAttachments* attachments = nullptr;  // Non-owning.

    ~GraphExecBox() noexcept {
        if (resource) {
            GILReleaseGuard gil;
            p_cuGraphExecDestroy(resource);
        }
        // The accumulator fields may be dangling after exec destruction.
        retry_deferred_cleanup();
    }
};

GraphExecBox* get_exec_box(const GraphExecHandle& h) noexcept {
    return const_cast<GraphExecBox*>(
        reinterpret_cast<const GraphExecBox*>(h.get()));
}

GraphExecHandle make_graph_exec_handle(
        CUgraphExec graph_exec, ExecAttachments* attachments) {
    struct RawGraphExecGuard {
        CUgraphExec resource;

        ~RawGraphExecGuard() noexcept {
            if (resource) {
                GILReleaseGuard gil;
                p_cuGraphExecDestroy(resource);
            }
            retry_deferred_cleanup();
        }
    } guard{graph_exec};

    auto box = std::make_shared<GraphExecBox>();
    box->resource = graph_exec;
    box->attachments = attachments;
    guard.resource = nullptr;
    return GraphExecHandle(box, &box->resource);
}

// Holds a fresh accumulator retained on the source graph across a CUDA call
// that propagates user objects into an exec. Releasing drops the source's
// reference: after successful propagation the exec keeps the accumulator
// alive, and otherwise this drops its last reference.
struct ExecAttachmentStaging {
    GraphHandle h_source;
    ExecAttachments* accumulator = nullptr;

    ~ExecAttachmentStaging() noexcept {
        release();
    }

    CUresult release() noexcept {
        if (!h_source || !accumulator) {
            return CUDA_SUCCESS;
        }
        const CUuserObject object = accumulator->object;
        const GraphHandle source = std::move(h_source);
        accumulator = nullptr;
        GILReleaseGuard gil;
        return p_cuGraphReleaseUserObject(*source, object, 1);
    }
};

// Create an accumulator and retain it on h_source, so that a following
// instantiation or whole-graph update propagates a reference into the exec.
CUresult stage_exec_attachments(
        const GraphHandle& h_source, ExecAttachmentStaging* out_staging) {
    if (!p_cuUserObjectCreate || !p_cuUserObjectRelease ||
        !p_cuGraphRetainUserObject || !p_cuGraphReleaseUserObject) {
        return CUDA_ERROR_NOT_SUPPORTED;
    }

    ensure_deferred_cleanup_ready();
    auto* accumulator = new ExecAttachments;

    CUuserObject object = nullptr;
    CUresult status;
    {
        GILReleaseGuard gil;
        status = p_cuUserObjectCreate(
            &object,
            static_cast<DeferredCleanupItem*>(accumulator),
            reinterpret_cast<CUhostFn>(enqueue_cleanup),
            1,
            CU_USER_OBJECT_NO_DESTRUCTOR_SYNC);
        if (status != CUDA_SUCCESS) {
            delete accumulator;
            return status;
        }
        accumulator->object = object;
        status = p_cuGraphRetainUserObject(
            *h_source, object, 1, CU_GRAPH_USER_OBJECT_MOVE);
        if (status != CUDA_SUCCESS) {
            // Dropping the last reference retires the accumulator.
            p_cuUserObjectRelease(object, 1);
            return status;
        }
    }

    out_staging->h_source = h_source;
    out_staging->accumulator = accumulator;
    return CUDA_SUCCESS;
}

}  // namespace

// State held by PreparedExecAttachment between preparation and commit. It keeps
// the exec alive and remembers the accumulator size before the append, so that
// rollback can drop owners staged for a mutation that CUDA rejected.
struct PreparedExecAttachmentState {
    GraphExecHandle h_exec;
    ExecAttachments* attachments = nullptr;
    size_t original_size = 0;

    PreparedExecAttachmentState(
            GraphExecHandle h_exec_,
            ExecAttachments* attachments_,
            size_t original_size_)
        : h_exec(std::move(h_exec_)),
          attachments(attachments_),
          original_size(original_size_) {}
};

void rollback_prepared_exec_attachment(
        PreparedExecAttachmentState* state) noexcept {
    if (!state) {
        return;
    }
    if (state->attachments) {
        while (state->attachments->owners.size() > state->original_size) {
            state->attachments->owners.pop_back();
        }
    }
    delete state;
}

GraphExecHandle create_graph_exec_handle(
        const GraphHandle& h_source,
        CUDA_GRAPH_INSTANTIATE_PARAMS* params) {
    if (!h_source || !*h_source || !params) {
        err = CUDA_ERROR_INVALID_VALUE;
        return {};
    }
    if (!p_cuGraphInstantiateWithParams) {
        err = CUDA_ERROR_NOT_SUPPORTED;
        return {};
    }

    ExecAttachmentStaging staging;
    if (CUDA_SUCCESS != (err = stage_exec_attachments(h_source, &staging))) {
        return {};
    }

    CUgraphExec graph_exec = nullptr;
    {
        GILReleaseGuard gil;
        err = p_cuGraphInstantiateWithParams(&graph_exec, *h_source, params);
    }
    if (err != CUDA_SUCCESS) {
        return {};
    }
    // CUDA can report a specific failure while returning success. The exec is
    // then unusable, so it stays unadopted for the caller to diagnose from
    // params->result_out.
    if (params->result_out != CUDA_GRAPH_INSTANTIATE_SUCCESS) {
        return {};
    }
    if (!graph_exec) {
        err = CUDA_ERROR_INVALID_VALUE;
        return {};
    }

    GraphExecHandle h_exec = make_graph_exec_handle(
        graph_exec, staging.accumulator);
    if (CUDA_SUCCESS != (err = staging.release())) {
        return {};
    }
    return h_exec;
}

CUresult graph_exec_update(
        const GraphExecHandle& h_exec,
        const GraphHandle& h_source,
        CUgraphExecUpdateResultInfo* result_info) {
    if (!h_exec || !h_source || !*h_source || !result_info) {
        return CUDA_ERROR_INVALID_VALUE;
    }
    if (!p_cuGraphExecUpdate) {
        return CUDA_ERROR_NOT_SUPPORTED;
    }

    GraphExecBox* box = get_exec_box(h_exec);
    if (!box->resource) {
        return CUDA_ERROR_INVALID_VALUE;
    }

    ExecAttachmentStaging staging;
    CUresult status = stage_exec_attachments(h_source, &staging);
    if (status != CUDA_SUCCESS) {
        return status;
    }

    {
        GILReleaseGuard gil;
        status = p_cuGraphExecUpdate(box->resource, *h_source, result_info);
    }
    if (status != CUDA_SUCCESS) {
        return status;
    }

    // CUDA may already have retired the old accumulator. Publish the new one
    // before releasing the source graph's temporary reference.
    box->attachments = staging.accumulator;
    return staging.release();
}

CUresult graph_prepare_exec_attachment(
        const GraphExecHandle& h_exec,
        OpaqueHandle owner0,
        OpaqueHandle owner1,
        PreparedExecAttachment* out_prepared) {
    if (!out_prepared) {
        return CUDA_ERROR_INVALID_VALUE;
    }
    out_prepared->reset();
    if (!h_exec) {
        return CUDA_ERROR_INVALID_VALUE;
    }

    GraphExecBox* box = get_exec_box(h_exec);
    if (!box->resource || !box->attachments) {
        return CUDA_ERROR_INVALID_VALUE;
    }

    ExecAttachments* attachments = box->attachments;
    const size_t original_size = attachments->owners.size();
    const size_t additions =
        static_cast<size_t>(static_cast<bool>(owner0)) +
        static_cast<size_t>(static_cast<bool>(owner1));
    // Reserve before staging so that rollback and commit cannot allocate.
    attachments->owners.reserve(original_size + additions);
    PreparedExecAttachment prepared(
        new PreparedExecAttachmentState(h_exec, attachments, original_size),
        PreparedExecAttachmentDeleter{rollback_prepared_exec_attachment});
    if (owner0) {
        attachments->owners.emplace_back(std::move(owner0));
    }
    if (owner1) {
        attachments->owners.emplace_back(std::move(owner1));
    }
    *out_prepared = std::move(prepared);
    return CUDA_SUCCESS;
}

void graph_commit_exec_attachment(
        PreparedExecAttachment& prepared) noexcept {
    delete prepared.release();
}

namespace {
struct GraphNodeBox {
    mutable CUgraphNode resource;
    GraphHandle h_graph;
};
}  // namespace

static const GraphNodeBox* get_box(const GraphNodeHandle& h) {
    const CUgraphNode* p = h.get();
    return reinterpret_cast<const GraphNodeBox*>(
        reinterpret_cast<const char*>(p) - offsetof(GraphNodeBox, resource)
    );
}

// graphs is ordered parent-before-child. Nulling a selected box marks its
// later descendants, whose parent pointers remain valid after list splicing.
// This permits one allocation-free sweep of the hierarchy.
void invalidate_child_graph_state(
        const GraphHandle& h_parent,
        CUgraphNode owner_node) noexcept {
    if (!h_parent || !owner_node) {
        return;
    }

    GraphBox* parent = get_box(h_parent);
    if (!parent->resource) {
        return;
    }
    GraphHierarchy& hierarchy = *parent->hierarchy;
    for (auto it = hierarchy.graphs.begin();
         it != hierarchy.graphs.end();) {
        auto graph = it++;
        bool is_owned_root = graph->parent == parent &&
                             graph->owner_node == owner_node;
        bool is_descendant = graph->parent &&
                             !graph->parent->resource;
        if (!is_owned_root && !is_descendant) {
            continue;
        }

        // Empty node_handles and invalidate each one.
        for (auto& entry : graph->node_handles.drain()) {
            if (GraphNodeHandle h_node = entry.second.lock()) {
                get_box(h_node)->resource = nullptr;
            }
        }
        graph_registry.unregister_handle(graph->resource);
        graph->resource = nullptr;
        graph->attachments.clear();
        hierarchy.graveyard.splice(
            hierarchy.graveyard.end(), hierarchy.graphs, graph);
    }
}

GraphNodeHandle create_graph_node_handle(CUgraphNode node, const GraphHandle& h_graph) {
    if (!node) {
        auto box = std::make_shared<const GraphNodeBox>(
            GraphNodeBox{nullptr, h_graph});
        return GraphNodeHandle(box, &box->resource);
    }

    GraphBox* graph = get_box(h_graph);
    return graph->node_handles.get_or_create(
        node,
        [node, &h_graph] {
            auto box = std::make_shared<const GraphNodeBox>(
                GraphNodeBox{node, h_graph});
            return GraphNodeHandle(box, &box->resource);
        });
}

GraphHandle graph_node_get_graph(const GraphNodeHandle& h) noexcept {
    return h ? get_box(h)->h_graph : GraphHandle{};
}

void invalidate_graph_node(const GraphNodeHandle& h) noexcept {
    if (!h) {
        return;
    }

    const GraphNodeBox* node_box = get_box(h);
    CUgraphNode node = node_box->resource;
    if (!node) {
        return;
    }
    GraphBox* graph = get_box(node_box->h_graph);
    graph->node_handles.unregister_handle(node);
    node_box->resource = nullptr;
}

// ============================================================================
// Graphics Resource Handles
// ============================================================================

namespace {
struct GraphicsResourceBox {
    CUgraphicsResource resource;
};
}  // namespace

GraphicsResourceHandle create_graphics_resource_handle(CUgraphicsResource resource) {
    auto box = std::shared_ptr<const GraphicsResourceBox>(
        new GraphicsResourceBox{resource},
        [](const GraphicsResourceBox* b) {
            GILReleaseGuard gil;
            p_cuGraphicsUnregisterResource(b->resource);
            delete b;
        }
    );
    return GraphicsResourceHandle(box, &box->resource);
}

// ============================================================================
// NVRTC Program Handles
// ============================================================================

namespace {
struct NvrtcProgramBox {
    nvrtcProgram resource;
};
}  // namespace

NvrtcProgramHandle create_nvrtc_program_handle(nvrtcProgram prog) {
    auto box = std::shared_ptr<NvrtcProgramBox>(
        new NvrtcProgramBox{prog},
        [](NvrtcProgramBox* b) {
            // Note: nvrtcDestroyProgram takes nvrtcProgram* and nulls it,
            // but we're deleting the box anyway so nulling is harmless.
            // Errors are ignored (standard destructor practice).
            p_nvrtcDestroyProgram(&b->resource);
            delete b;
        }
    );
    return NvrtcProgramHandle(box, &box->resource);
}

NvrtcProgramHandle create_nvrtc_program_handle_ref(nvrtcProgram prog) {
    auto box = std::make_shared<NvrtcProgramBox>(NvrtcProgramBox{prog});
    return NvrtcProgramHandle(box, &box->resource);
}

// ============================================================================
// NVVM Program Handles
// ============================================================================

namespace {
struct NvvmProgramBox {
    NvvmProgramValue resource;
};
}  // namespace

NvvmProgramHandle create_nvvm_program_handle(nvvmProgram prog) {
    auto box = std::shared_ptr<NvvmProgramBox>(
        new NvvmProgramBox{{prog}},
        [](NvvmProgramBox* b) {
            // Note: nvvmDestroyProgram takes nvvmProgram* and nulls it,
            // but we're deleting the box anyway so nulling is harmless.
            // If NVVM is not available, the function pointer is null.
            if (p_nvvmDestroyProgram) {
                p_nvvmDestroyProgram(&b->resource.raw);
            }
            delete b;
        }
    );
    return NvvmProgramHandle(box, &box->resource);
}

NvvmProgramHandle create_nvvm_program_handle_ref(nvvmProgram prog) {
    auto box = std::make_shared<NvvmProgramBox>(NvvmProgramBox{{prog}});
    return NvvmProgramHandle(box, &box->resource);
}

// ============================================================================
// nvJitLink Handles
// ============================================================================

namespace {
struct NvJitLinkBox {
    NvJitLinkValue resource;
};
}  // namespace

NvJitLinkHandle create_nvjitlink_handle(nvJitLink_t handle) {
    auto box = std::shared_ptr<NvJitLinkBox>(
        new NvJitLinkBox{{handle}},
        [](NvJitLinkBox* b) {
            // Note: nvJitLinkDestroy takes nvJitLinkHandle* and nulls it,
            // but we're deleting the box anyway so nulling is harmless.
            // If nvJitLink is not available, the function pointer is null.
            if (p_nvJitLinkDestroy) {
                p_nvJitLinkDestroy(&b->resource.raw);
            }
            delete b;
        }
    );
    return NvJitLinkHandle(box, &box->resource);
}

NvJitLinkHandle create_nvjitlink_handle_ref(nvJitLink_t handle) {
    auto box = std::make_shared<NvJitLinkBox>(NvJitLinkBox{{handle}});
    return NvJitLinkHandle(box, &box->resource);
}

// ============================================================================
// cuLink Handles
// ============================================================================

namespace {
struct CuLinkBox {
    CUlinkState resource;
};
}  // namespace

CuLinkHandle create_culink_handle(CUlinkState state) {
    auto box = std::shared_ptr<CuLinkBox>(
        new CuLinkBox{state},
        [](CuLinkBox* b) {
            // cuLinkDestroy takes CUlinkState by value (not pointer).
            // Errors are ignored (standard destructor practice).
            if (p_cuLinkDestroy) {
                p_cuLinkDestroy(b->resource);
            }
            delete b;
        }
    );
    return CuLinkHandle(box, &box->resource);
}

CuLinkHandle create_culink_handle_ref(CUlinkState state) {
    auto box = std::make_shared<CuLinkBox>(CuLinkBox{state});
    return CuLinkHandle(box, &box->resource);
}

// ============================================================================
// File Descriptor Handles
// ============================================================================

FileDescriptorHandle create_fd_handle(int fd) {
#ifdef _WIN32
    throw std::runtime_error("create_fd_handle is not supported on Windows");
#else
    return FileDescriptorHandle(
        new int(fd),
        [](const int* p) { ::close(*p); delete p; }
    );
#endif
}

FileDescriptorHandle create_fd_handle_ref(int fd) {
#ifdef _WIN32
    throw std::runtime_error("create_fd_handle_ref is not supported on Windows");
#else
    return std::make_shared<const int>(fd);
#endif
}

// ============================================================================
// Array / mipmapped-array / texture / surface handles (PR #467)
// ============================================================================

namespace {
struct ArrayBox {
    CUarray resource;
    // Non-null only for a mipmap-level view: keeps the parent mipmap (the real
    // owner of the level's storage) alive for as long as the level is held.
    MipmappedArrayHandle h_parent;
};

struct MipmappedArrayBox {
    CUmipmappedArray resource;
};

struct TexObjectBox {
    // Tagged so TexObjectHandle is a distinct C++ type from DevicePtrHandle /
    // SurfObjectHandle (all wrap `unsigned long long`).
    TexObjectValue resource;
    // Type-erased backing dependency (OpaqueArrayHandle / MipmappedArrayHandle /
    // DevicePtrHandle). The texture's resource is a union; we only need to keep
    // whichever backing it was built from alive, never to dereference it.
    std::shared_ptr<const void> h_backing;
};

struct SurfObjectBox {
    SurfObjectValue resource;
    OpaqueArrayHandle h_array;  // surfaces are always array-backed
};
}  // namespace

OpaqueArrayHandle create_array_handle(const CUDA_ARRAY3D_DESCRIPTOR& desc) {
    GILReleaseGuard gil;
    CUarray arr;
    if (CUDA_SUCCESS != (err = p_cuArray3DCreate(&arr, &desc))) {
        return {};
    }
    // Allocation and adoption share the same owning lifetime; the only
    // difference is who calls cuArray3DCreate. Delegate so the owning box and
    // its destroy-on-last-ref deleter are defined in exactly one place.
    return create_array_handle_owning(arr);
}

OpaqueArrayHandle create_array_handle_ref(CUarray arr) {
    if (!arr) {
        return {};
    }
    auto box = std::make_shared<const ArrayBox>(ArrayBox{arr, {}});
    return OpaqueArrayHandle(box, &box->resource);
}

OpaqueArrayHandle create_array_handle_owning(CUarray arr) {
    if (!arr) {
        return {};
    }
    auto box = std::shared_ptr<const ArrayBox>(
        new ArrayBox{arr, {}},
        [](const ArrayBox* b) {
            GILReleaseGuard gil;
            p_cuArrayDestroy(b->resource);
            delete b;
        }
    );
    return OpaqueArrayHandle(box, &box->resource);
}

OpaqueArrayHandle create_array_level_handle(const MipmappedArrayHandle& h_mip, unsigned int level) {
    GILReleaseGuard gil;
    CUarray arr;
    if (CUDA_SUCCESS != (err = p_cuMipmappedArrayGetLevel(&arr, as_cu(h_mip), level))) {
        return {};
    }
    // Non-owning level view: storage belongs to the mipmap. Embed the mipmap
    // handle so the parent outlives this level; the deleter does not destroy.
    auto box = std::shared_ptr<const ArrayBox>(
        new ArrayBox{arr, h_mip},
        [](const ArrayBox* b) { delete b; }
    );
    return OpaqueArrayHandle(box, &box->resource);
}

MipmappedArrayHandle create_mipmapped_array_handle(const CUDA_ARRAY3D_DESCRIPTOR& desc,
                                                   unsigned int num_levels) {
    GILReleaseGuard gil;
    CUmipmappedArray mip;
    if (CUDA_SUCCESS != (err = p_cuMipmappedArrayCreate(&mip, &desc, num_levels))) {
        return {};
    }
    auto box = std::shared_ptr<const MipmappedArrayBox>(
        new MipmappedArrayBox{mip},
        [](const MipmappedArrayBox* b) {
            GILReleaseGuard gil;
            p_cuMipmappedArrayDestroy(b->resource);
            delete b;
        }
    );
    return MipmappedArrayHandle(box, &box->resource);
}

namespace {
TexObjectHandle make_tex_object_handle(const CUDA_RESOURCE_DESC& res,
                                       const CUDA_TEXTURE_DESC& tex,
                                       std::shared_ptr<const void> h_backing) {
    GILReleaseGuard gil;
    CUtexObject obj;
    if (CUDA_SUCCESS != (err = p_cuTexObjectCreate(&obj, &res, &tex, nullptr))) {
        return {};
    }
    auto box = std::shared_ptr<const TexObjectBox>(
        new TexObjectBox{TexObjectValue{obj}, std::move(h_backing)},
        [](const TexObjectBox* b) {
            GILReleaseGuard gil;
            p_cuTexObjectDestroy(b->resource.raw);
            delete b;
        }
    );
    return TexObjectHandle(box, &box->resource);
}
}  // namespace

TexObjectHandle create_tex_object_handle_array(const CUDA_RESOURCE_DESC& res,
                                               const CUDA_TEXTURE_DESC& tex,
                                               const OpaqueArrayHandle& h_backing) {
    return make_tex_object_handle(res, tex, h_backing);
}

TexObjectHandle create_tex_object_handle_mipmap(const CUDA_RESOURCE_DESC& res,
                                                const CUDA_TEXTURE_DESC& tex,
                                                const MipmappedArrayHandle& h_backing) {
    return make_tex_object_handle(res, tex, h_backing);
}

TexObjectHandle create_tex_object_handle_linear(const CUDA_RESOURCE_DESC& res,
                                                const CUDA_TEXTURE_DESC& tex,
                                                const DevicePtrHandle& h_backing) {
    return make_tex_object_handle(res, tex, h_backing);
}

SurfObjectHandle create_surf_object_handle(const CUDA_RESOURCE_DESC& res,
                                           const OpaqueArrayHandle& h_backing) {
    GILReleaseGuard gil;
    CUsurfObject obj;
    if (CUDA_SUCCESS != (err = p_cuSurfObjectCreate(&obj, &res))) {
        return {};
    }
    auto box = std::shared_ptr<const SurfObjectBox>(
        new SurfObjectBox{SurfObjectValue{obj}, h_backing},
        [](const SurfObjectBox* b) {
            GILReleaseGuard gil;
            p_cuSurfObjectDestroy(b->resource.raw);
            delete b;
        }
    );
    return SurfObjectHandle(box, &box->resource);
}

// ============================================================================
// SM resource split wrapper
// ============================================================================

CUresult sm_resource_split(CUdevResource* result, unsigned int nbGroups,
                           const CUdevResource* input, CUdevResource* remainder,
                           unsigned int flags, void* groupParams) {
#if CUDA_VERSION >= 13010
    if (!p_cuDevSmResourceSplit) {
        return CUDA_ERROR_NOT_SUPPORTED;
    }
    return p_cuDevSmResourceSplit(
        result, nbGroups, input, remainder, flags,
        static_cast<CU_DEV_SM_RESOURCE_GROUP_PARAMS*>(groupParams));
#else
    return CUDA_ERROR_NOT_SUPPORTED;
#endif
}

bool has_sm_resource_split() noexcept {
    return p_cuDevSmResourceSplit != nullptr;
}

}  // namespace cuda_core
