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

import ctypes
import sys

from cuda.bindings import driver

try:
    import numpy as np
except ImportError:
    np = None
import platform
import re

import pytest
from conftest import (
    create_managed_memory_resource_or_skip,
    create_pinned_memory_resource_or_xfail,
    skip_if_managed_memory_unsupported,
    skip_if_pinned_memory_unsupported,
)
from helpers import supports_ipc_mempool
from helpers.buffers import DummyDeviceMemoryResource, DummyUnifiedMemoryResource, TrackingMR
from helpers.constants import POOL_SIZE

from cuda.core import (
    Buffer,
    Device,
    DeviceMemoryResource,
    DeviceMemoryResourceOptions,
    GraphMemoryResource,
    LegacyPinnedMemoryResource,
    ManagedMemoryResource,
    ManagedMemoryResourceOptions,
    MemoryResource,
    PinnedMemoryResource,
    PinnedMemoryResourceOptions,
    VirtualMemoryResource,
    VirtualMemoryResourceOptions,
)
from cuda.core import (
    system as ccx_system,
)
from cuda.core._dlpack import DLDeviceType
from cuda.core._memory._ipc import IPCBufferDescriptor
from cuda.core._utils.cuda_utils import CUDAError, handle_return
from cuda.core.typing import (
    ManagedMemoryLocationType,
    VirtualMemoryAccessType,
    VirtualMemoryAllocationType,
    VirtualMemoryGranularityType,
    VirtualMemoryHandleType,
    VirtualMemoryLocationType,
)
from cuda.core.utils import StridedMemoryView
from cuda_python_test_helpers import IS_WINDOWS


def _allocate_pinned_buffer_or_xfail(mr, size, *, device):
    try:
        return mr.allocate(size, stream=device.default_stream)
    except CUDAError as exc:
        if "CUDA_ERROR_OUT_OF_MEMORY" in str(exc):
            pytest.xfail("TODO(#9999): Resolve CUDA_ERROR_OUT_OF_MEMORY")
        raise
    except RuntimeError as exc:
        if "Failed to allocate memory from pool" in str(exc):
            pytest.xfail("TODO(#9999): Resolve Failed to allocate memory from pool")
        raise


class DummyHostMemoryResource(MemoryResource):
    # Pure-host ctypes allocation; stream is accepted for interface
    # conformance but ignored.
    def __init__(self):
        pass

    def allocate(self, size, *, stream=None) -> Buffer:
        # Allocate a ctypes buffer of size `size`
        ptr = (ctypes.c_byte * size)()
        self._ptr = ptr
        return Buffer.from_handle(ptr=ctypes.addressof(ptr), size=size, mr=self)

    def deallocate(self, ptr, size, *, stream=None):
        del self._ptr

    @property
    def is_device_accessible(self) -> bool:
        return False

    @property
    def is_host_accessible(self) -> bool:
        return True

    @property
    def device_id(self) -> int:
        raise RuntimeError("the pinned memory resource is not bound to any GPU")


class DummyPinnedMemoryResource(MemoryResource):
    # cuMemAllocHost / cuMemFreeHost are synchronous; stream is accepted
    # for interface conformance but ignored.
    def __init__(self, device):
        self.device = device

    def allocate(self, size, *, stream=None) -> Buffer:
        ptr = handle_return(driver.cuMemAllocHost(size))
        return Buffer.from_handle(ptr=ptr, size=size, mr=self)

    def deallocate(self, ptr, size, *, stream=None):
        handle_return(driver.cuMemFreeHost(ptr))

    @property
    def is_device_accessible(self) -> bool:
        return True

    @property
    def is_host_accessible(self) -> bool:
        return True

    @property
    def device_id(self) -> int:
        raise RuntimeError("the pinned memory resource is not bound to any GPU")


class NullMemoryResource(DummyHostMemoryResource):
    @property
    def is_host_accessible(self) -> bool:
        return False


def test_package_contents():
    expected = [
        "Buffer",
        "DeviceMemoryResource",
        "DeviceMemoryResourceOptions",
        "GraphMemoryResource",
        "LegacyPinnedMemoryResource",
        "ManagedBuffer",
        "ManagedMemoryResource",
        "ManagedMemoryResourceOptions",
        "MemoryResource",
        "PinnedMemoryResource",
        "PinnedMemoryResourceOptions",
        "VirtualMemoryResource",
        "VirtualMemoryResourceOptions",
    ]
    d = {}
    exec("from cuda.core._memory import *", d)  # noqa: S102
    d = {k: v for k, v in d.items() if not k.startswith("__")}
    assert sorted(expected) == sorted(d.keys())


def buffer_initialization(dummy_mr: MemoryResource):
    buffer = dummy_mr.allocate(size=1024)
    assert buffer.handle != 0
    assert buffer.size == 1024
    assert buffer.memory_resource == dummy_mr
    assert buffer.is_device_accessible == dummy_mr.is_device_accessible
    assert buffer.is_host_accessible == dummy_mr.is_host_accessible
    assert not buffer.is_mapped
    buffer.close()


def test_buffer_initialization():
    device = Device()
    device.set_current()
    buffer_initialization(DummyDeviceMemoryResource(device))
    buffer_initialization(DummyHostMemoryResource())
    buffer_initialization(DummyUnifiedMemoryResource(device))
    buffer_initialization(DummyPinnedMemoryResource(device))
    with pytest.raises(TypeError):
        buffer_initialization(MemoryResource())


def buffer_copy_to(dummy_mr: MemoryResource, device: Device, check=False):
    src_buffer = dummy_mr.allocate(size=1024)
    dst_buffer = dummy_mr.allocate(size=1024)
    stream = device.create_stream()

    if check:
        src_ptr = ctypes.cast(src_buffer.handle, ctypes.POINTER(ctypes.c_byte))
        for i in range(1024):
            src_ptr[i] = ctypes.c_byte(i)

    src_buffer.copy_to(dst_buffer, stream=stream)
    device.sync()

    if check:
        dst_ptr = ctypes.cast(dst_buffer.handle, ctypes.POINTER(ctypes.c_byte))

        for i in range(10):
            assert dst_ptr[i] == src_ptr[i]

    dst_buffer.close()
    src_buffer.close()


def test_buffer_copy_to():
    device = Device()
    device.set_current()
    buffer_copy_to(DummyDeviceMemoryResource(device), device)
    buffer_copy_to(DummyUnifiedMemoryResource(device), device)
    buffer_copy_to(DummyPinnedMemoryResource(device), device, check=True)


def buffer_copy_from(dummy_mr: MemoryResource, device, check=False):
    src_buffer = dummy_mr.allocate(size=1024)
    dst_buffer = dummy_mr.allocate(size=1024)
    stream = device.create_stream()

    if check:
        src_ptr = ctypes.cast(src_buffer.handle, ctypes.POINTER(ctypes.c_byte))
        for i in range(1024):
            src_ptr[i] = ctypes.c_byte(i)

    dst_buffer.copy_from(src_buffer, stream=stream)
    device.sync()

    if check:
        dst_ptr = ctypes.cast(dst_buffer.handle, ctypes.POINTER(ctypes.c_byte))

        for i in range(10):
            assert dst_ptr[i] == src_ptr[i]

    dst_buffer.close()
    src_buffer.close()


def test_buffer_copy_from():
    device = Device()
    device.set_current()
    buffer_copy_from(DummyDeviceMemoryResource(device), device)
    buffer_copy_from(DummyUnifiedMemoryResource(device), device)
    buffer_copy_from(DummyPinnedMemoryResource(device), device, check=True)


def _bytes_repeat(pattern: bytes, size: int) -> bytes:
    assert len(pattern) > 0
    assert size % len(pattern) == 0
    return pattern * (size // len(pattern))


def _pattern_bytes(value) -> bytes:
    if isinstance(value, int):
        return bytes([value])
    return bytes(memoryview(value).cast("B"))


@pytest.fixture(params=["device", "unified", "pinned"])
def fill_env(request):
    device = Device()
    device.set_current()
    if request.param == "device":
        mr = DummyDeviceMemoryResource(device)
    elif request.param == "unified":
        mr = DummyUnifiedMemoryResource(device)
    else:
        mr = DummyPinnedMemoryResource(device)
    return device, mr


_FILL_SIZE = 64  # Keep small; divisible by 1/2/4.

_FILL_CASES = [
    # int -> 1-byte pattern
    pytest.param(0x42, _FILL_SIZE, None, id="int-0x42"),
    pytest.param(-1, _FILL_SIZE, OverflowError, id="int-neg"),
    pytest.param(256, _FILL_SIZE, OverflowError, id="int-256"),
    pytest.param(1000, _FILL_SIZE, OverflowError, id="int-1000"),
    # bad type
    pytest.param("invalid", _FILL_SIZE, TypeError, id="bad-type-str"),
    # bytes-like patterns
    pytest.param(b"\x7f", _FILL_SIZE, None, id="bytes-1"),
    pytest.param(b"\x34\x12", _FILL_SIZE, None, id="bytes-2"),
    pytest.param(b"\xef\xbe\xad\xde", _FILL_SIZE, None, id="bytes-4"),
    pytest.param(b"\x34\x12", _FILL_SIZE + 1, ValueError, id="bytes-2-bad-size"),
    pytest.param(b"\xef\xbe\xad\xde", _FILL_SIZE + 2, ValueError, id="bytes-4-bad-size"),
    pytest.param(b"", _FILL_SIZE, ValueError, id="bytes-0"),
    pytest.param(b"\x01\x02\x03", _FILL_SIZE, ValueError, id="bytes-3"),
]

if np is not None:
    _FILL_CASES.extend(
        [
            # 8-bit patterns
            pytest.param(np.uint8(0), _FILL_SIZE, None, id="np-uint8-0"),
            pytest.param(np.uint8(255), _FILL_SIZE, None, id="np-uint8-255"),
            pytest.param(np.int8(-1), _FILL_SIZE, None, id="np-int8--1"),
            pytest.param(np.int8(127), _FILL_SIZE, None, id="np-int8-127"),
            pytest.param(np.int8(-128), _FILL_SIZE, None, id="np-int8--128"),
            # 16-bit patterns
            pytest.param(np.uint16(0x1234), _FILL_SIZE, None, id="np-uint16-0x1234"),
            pytest.param(np.uint16(0xFFFF), _FILL_SIZE, None, id="np-uint16-0xFFFF"),
            pytest.param(np.int16(-1), _FILL_SIZE, None, id="np-int16--1"),
            pytest.param(np.int16(32767), _FILL_SIZE, None, id="np-int16-max"),
            pytest.param(np.int16(-32768), _FILL_SIZE, None, id="np-int16-min"),
            pytest.param(np.uint16(0x1234), _FILL_SIZE + 1, ValueError, id="np-uint16-bad-size"),
            # 32-bit patterns
            pytest.param(np.uint32(0xDEADBEEF), _FILL_SIZE, None, id="np-uint32-0xDEADBEEF"),
            pytest.param(np.uint32(0xFFFFFFFF), _FILL_SIZE, None, id="np-uint32-0xFFFFFFFF"),
            pytest.param(np.int32(-1), _FILL_SIZE, None, id="np-int32--1"),
            pytest.param(np.int32(2147483647), _FILL_SIZE, None, id="np-int32-max"),
            pytest.param(np.int32(-2147483648), _FILL_SIZE, None, id="np-int32-min"),
            pytest.param(np.uint32(0xDEADBEEF), _FILL_SIZE + 2, ValueError, id="np-uint32-bad-size"),
            # float32 (bit-pattern fill)
            pytest.param(np.float32(1.0), _FILL_SIZE, None, id="np-float32-1.0"),
            # 64-bit patterns should error (8-byte pattern)
            pytest.param(np.uint64(0), _FILL_SIZE, ValueError, id="np-uint64-err"),
            pytest.param(np.int64(0), _FILL_SIZE, ValueError, id="np-int64-err"),
            pytest.param(np.float64(0), _FILL_SIZE, ValueError, id="np-float64-err"),
        ]
    )


@pytest.mark.parametrize("value,size,exc", _FILL_CASES)
def test_buffer_fill(fill_env, value, size, exc):
    device, mr = fill_env
    stream = device.create_stream()
    buffer = mr.allocate(size=size)
    try:
        if exc is not None:
            with pytest.raises(exc):
                buffer.fill(value, stream=stream)
            return

        buffer.fill(value, stream=stream)
        device.sync()

        # Verify contents only for host-accessible buffers.
        if buffer.is_host_accessible:
            pat = _pattern_bytes(value)
            got = ctypes.string_at(int(buffer.handle), size)
            assert got == _bytes_repeat(pat, size)
    finally:
        buffer.close()


def buffer_close(dummy_mr: MemoryResource):
    buffer = dummy_mr.allocate(size=1024)
    buffer.close()
    assert buffer.handle == 0
    assert buffer.memory_resource is None


def test_buffer_close():
    device = Device()
    device.set_current()
    buffer_close(DummyDeviceMemoryResource(device))
    buffer_close(DummyHostMemoryResource())
    buffer_close(DummyUnifiedMemoryResource(device))
    buffer_close(DummyPinnedMemoryResource(device))


def test_buffer_external_host():
    a = (ctypes.c_byte * 20)()
    ptr = ctypes.addressof(a)
    buffer = Buffer.from_handle(ptr, 20, owner=a)
    assert not buffer.is_device_accessible
    assert buffer.is_host_accessible
    assert buffer.device_id == -1
    buffer.close()


@pytest.mark.parametrize("change_device", [True, False])
def test_buffer_external_device(change_device):
    n = ccx_system.get_num_devices()
    if n < 1:
        pytest.skip("No devices found")
    dev_id = n - 1
    d = Device(dev_id)
    d.set_current()
    buffer_ = d.allocate(size=32, stream=d.default_stream)

    if change_device:
        # let's switch to a different device if possibe
        # to make sure we get the original device id
        d = Device(0)
        d.set_current()

    buffer = Buffer.from_handle(int(buffer_.handle), 32)
    assert buffer.is_device_accessible
    assert not buffer.is_host_accessible
    assert buffer.device_id == dev_id
    buffer.close()
    buffer_.close()


@pytest.mark.parametrize("change_device", [True, False])
def test_buffer_external_pinned_alloc(change_device):
    n = ccx_system.get_num_devices()
    if n < 1:
        pytest.skip("No devices found")
    dev_id = n - 1
    d = Device(dev_id)
    d.set_current()
    mr = DummyPinnedMemoryResource(d)
    buffer_ = mr.allocate(size=32)

    if change_device:
        # let's switch to a different device if possibe
        # to make sure we get the original device id
        d = Device(0)
        d.set_current()

    buffer = Buffer.from_handle(int(buffer_.handle), 32)
    assert buffer.is_device_accessible
    assert buffer.is_host_accessible
    assert buffer.device_id == dev_id
    buffer.close()
    buffer_.close()


@pytest.mark.parametrize("change_device", [True, False])
def test_buffer_external_pinned_registered(change_device):
    n = ccx_system.get_num_devices()
    if n < 1:
        pytest.skip("No devices found")
    dev_id = n - 1
    d = Device(dev_id)
    d.set_current()
    a = (ctypes.c_byte * 20)()
    ptr = ctypes.addressof(a)

    buffer = Buffer.from_handle(ptr, 20, owner=ptr)
    assert not buffer.is_device_accessible
    assert buffer.is_host_accessible
    assert buffer.device_id == -1

    handle_return(driver.cuMemHostRegister(ptr, 20, 0))
    try:
        if change_device:
            # let's switch to a different device if possibe
            # to make sure we get the original device id
            d = Device(0)
            d.set_current()

        buffer = Buffer.from_handle(ptr, 20, owner=ptr)
        assert buffer.is_device_accessible
        assert buffer.is_host_accessible
        assert buffer.device_id == dev_id
        buffer.close()
    finally:
        handle_return(driver.cuMemHostUnregister(ptr))


@pytest.mark.parametrize("change_device", [True, False])
def test_buffer_external_managed(change_device):
    n = ccx_system.get_num_devices()
    if n < 1:
        pytest.skip("No devices found")
    dev_id = n - 1
    d = Device(dev_id)
    d.set_current()
    ptr = None
    try:
        ptr = handle_return(driver.cuMemAllocManaged(32, driver.CUmemAttach_flags.CU_MEM_ATTACH_GLOBAL.value))
        if change_device:
            # let's switch to a different device if possibe
            # to make sure we get the original device id
            d = Device(0)
            d.set_current()
        buffer = Buffer.from_handle(ptr, 32)
        assert buffer.is_device_accessible
        assert buffer.is_host_accessible
        assert buffer.device_id == dev_id
    finally:
        if ptr is not None:
            handle_return(driver.cuMemFree(ptr))


def test_mr_deallocate_called_on_close():
    """Buffer.from_handle(mr=mr) calls mr.deallocate() on close (issue #1619)."""
    device = Device()
    device.set_current()
    mr = TrackingMR()
    buf = mr.allocate(1024)
    assert len(mr.active) == 1
    buf.close()
    assert len(mr.active) == 0


def test_mr_deallocate_called_on_gc():
    """Buffer.from_handle(mr=mr) calls mr.deallocate() on GC (issue #1619)."""
    import gc

    device = Device()
    device.set_current()
    mr = TrackingMR()
    buf = mr.allocate(1024)
    assert len(mr.active) == 1
    del buf
    gc.collect()
    assert len(mr.active) == 0


def test_mr_deallocate_receives_stream():
    """Buffer.close(stream) forwards the stream to mr.deallocate() (issue #1619)."""
    device = Device()
    device.set_current()
    stream = device.create_stream()
    received = {}

    class StreamCaptureMR(TrackingMR):
        def deallocate(self, ptr, size, *, stream=None):
            received["stream"] = stream
            super().deallocate(ptr, size, stream=stream)

    mr = StreamCaptureMR()
    buf = mr.allocate(1024)
    buf.close(stream)
    assert received["stream"].handle == stream.handle


def test_mr_dealloc_callback_falls_back_to_default_stream():
    """When a Buffer's device-pointer handle has no attached deallocation
    stream (e.g. buffers minted via :meth:`Buffer.from_handle` from DLPack
    import, IPC import, or third-party adapters), the C++ deleter callback
    must fall back to the default stream rather than passing ``stream=None``
    to ``mr.deallocate``. Stream-ordered MRs validate the stream and would
    otherwise raise ``TypeError`` from inside the ``noexcept`` callback,
    which only logs a warning and silently leaks the allocation. See
    `#2001 <https://github.com/NVIDIA/cuda-python/issues/2001>`__.
    """
    import gc

    from cuda.core._stream import Stream_accept, default_stream

    device = Device()
    device.set_current()
    captured = {}

    class StrictCapturingMR(MemoryResource):
        # Models a stream-ordered MR: deallocate validates the stream
        # the same way DeviceMemoryResource.deallocate does.
        @property
        def is_device_accessible(self):
            return True

        @property
        def is_host_accessible(self):
            return False

        @property
        def device_id(self):
            return device.device_id

        def allocate(self, size, *, stream):
            raise NotImplementedError  # not used; we use from_handle below

        def deallocate(self, ptr, size, *, stream):
            captured["stream"] = Stream_accept(stream)

    mr = StrictCapturingMR()
    # Buffer.from_handle binds mr but does not attach a deallocation stream.
    # ptr=1 is fine because StrictCapturingMR.deallocate does not free.
    buf = Buffer.from_handle(1, 1024, mr=mr)
    del buf
    gc.collect()

    assert "stream" in captured, "deallocate was not invoked (callback raised and leaked)"
    assert captured["stream"].handle == default_stream().handle


def test_memory_resource_and_owner_disallowed():
    with pytest.raises(ValueError, match="cannot be both specified together"):
        a = (ctypes.c_byte * 20)()
        ptr = ctypes.addressof(a)
        Buffer.from_handle(ptr, 20, mr=DummyDeviceMemoryResource(Device()), owner=a)


def test_owner_close():
    a = (ctypes.c_byte * 20)()
    ptr = ctypes.addressof(a)
    before = sys.getrefcount(a)
    buffer = Buffer.from_handle(ptr, 20, owner=a)
    assert sys.getrefcount(a) != before
    buffer.close()
    after = sys.getrefcount(a)
    assert after == before


def test_buffer_dunder_dlpack():
    device = Device()
    device.set_current()
    dummy_mr = DummyDeviceMemoryResource(device)
    buffer = dummy_mr.allocate(size=1024)
    capsule = buffer.__dlpack__()
    assert "dltensor" in repr(capsule)
    capsule = buffer.__dlpack__(max_version=(1, 0))
    assert "dltensor" in repr(capsule)
    with pytest.raises(BufferError, match=r"^Sorry, not supported: dl_device other than None$"):
        buffer.__dlpack__(dl_device=())
    with pytest.raises(BufferError, match=r"^Sorry, not supported: copy=True$"):
        buffer.__dlpack__(copy=True)
    with pytest.raises(BufferError, match=r"^Expected max_version tuple\[int, int\], got \(\)$"):
        buffer.__dlpack__(max_version=())
    with pytest.raises(BufferError, match=r"^Expected max_version tuple\[int, int\], got \(9, 8, 7\)$"):
        buffer.__dlpack__(max_version=(9, 8, 7))


@pytest.mark.parametrize(
    ("DummyMR", "expected"),
    [
        (DummyDeviceMemoryResource, (DLDeviceType.kDLCUDA, 0)),
        (DummyHostMemoryResource, (DLDeviceType.kDLCPU, 0)),
        (DummyUnifiedMemoryResource, (DLDeviceType.kDLCUDAManaged, 0)),
        (DummyPinnedMemoryResource, (DLDeviceType.kDLCUDAHost, 0)),
    ],
)
def test_buffer_dunder_dlpack_device_success(DummyMR, expected):
    device = Device()
    device.set_current()
    dummy_mr = DummyMR() if DummyMR is DummyHostMemoryResource else DummyMR(device)
    buffer = dummy_mr.allocate(size=1024)
    assert buffer.__dlpack_device__() == expected


def test_buffer_dunder_dlpack_device_failure():
    dummy_mr = NullMemoryResource()
    buffer = dummy_mr.allocate(size=1024)
    with pytest.raises(BufferError, match=r"^buffer is neither device-accessible nor host-accessible$"):
        buffer.__dlpack_device__()


def test_buffer_dlpack_failure_clean_up():
    dummy_mr = NullMemoryResource()
    buffer = dummy_mr.allocate(size=1024)
    before = sys.getrefcount(buffer)
    with pytest.raises(BufferError, match="buffer is neither device-accessible nor host-accessible"):
        buffer.__dlpack__()
    after = sys.getrefcount(buffer)
    # we use the buffer refcount as sentinel for proper clean-up here,
    # hoping that malloc and frees did the right thing
    # as they are handled by the same deleter
    assert after == before


def test_managed_buffer_dlpack_roundtrip_device_type():
    """Verify that a managed Buffer round-trips through DLPack with kDLCUDAManaged."""
    device = Device()
    device.set_current()
    skip_if_managed_memory_unsupported(device)
    mr = DummyUnifiedMemoryResource(device)
    buf = mr.allocate(size=1024)

    # Buffer-level classification should report managed.
    assert buf.__dlpack_device__() == (DLDeviceType.kDLCUDAManaged, 0)

    # The end-to-end path: Buffer -> DLPack capsule -> StridedMemoryView
    # must preserve kDLCUDAManaged rather than downgrading to kDLCUDAHost.
    view = StridedMemoryView.from_any_interface(buf, stream_ptr=-1)
    assert view.__dlpack_device__() == (int(DLDeviceType.kDLCUDAManaged), 0)


def test_managed_memory_resource_buffer_dlpack_device_type():
    """Verify that pool-allocated managed memory reports kDLCUDAManaged.

    Allocations from ManagedMemoryResource go through cuMemAllocFromPoolAsync
    on a CU_MEM_ALLOCATION_TYPE_MANAGED pool, which does not set
    CU_POINTER_ATTRIBUTE_IS_MANAGED.  The Buffer must still report itself as
    managed via its memory resource so that DLPack consumers (e.g. CCCL's
    make_tma_descriptor) accept the buffer.
    """
    device = Device()
    device.set_current()
    skip_if_managed_memory_unsupported(device)
    mr = create_managed_memory_resource_or_skip(ManagedMemoryResourceOptions(preferred_location=device.device_id))
    buf = mr.allocate(1024, stream=device.default_stream)

    assert mr.is_managed
    assert buf.is_managed
    assert buf.__dlpack_device__() == (DLDeviceType.kDLCUDAManaged, 0)

    view = StridedMemoryView.from_any_interface(buf, stream_ptr=-1)
    assert view.__dlpack_device__() == (int(DLDeviceType.kDLCUDAManaged), 0)


@pytest.mark.parametrize("mr_kind", ["device", "pinned"])
def test_non_managed_resources_report_not_managed(mr_kind):
    """Non-managed memory resources must report is_managed=False."""
    device = Device()
    device.set_current()
    if not device.properties.memory_pools_supported:
        pytest.skip("Device does not support mempool operations")
    if mr_kind == "device":
        mr = DeviceMemoryResource(device)
    else:
        skip_if_pinned_memory_unsupported(device)
        mr = create_pinned_memory_resource_or_xfail(xfail_device=device)
    assert mr.is_managed is False
    buf = (
        _allocate_pinned_buffer_or_xfail(mr, 1024, device=device)
        if mr_kind == "pinned"
        else mr.allocate(1024, stream=device.default_stream)
    )
    assert buf.is_managed is False
    buf.close()


@pytest.mark.parametrize("use_device_object", [True, False])
def test_device_memory_resource_initialization(use_device_object):
    """Test that DeviceMemoryResource can be initialized successfully.

    This test verifies that the DeviceMemoryResource initializes properly,
    including the release threshold configuration for performance optimization.
    """
    device = Device()

    if not device.properties.memory_pools_supported:
        pytest.skip("Device does not support mempool operations")

    device.set_current()

    # This should succeed and configure the memory pool release threshold.
    # The resource can be constructed from either a device or device ordinal.
    device_arg = device if use_device_object else device.device_id
    mr = DeviceMemoryResource(device_arg)

    # Verify basic properties
    assert mr.device_id == device.device_id
    assert mr.is_device_accessible
    assert not mr.is_host_accessible
    assert not mr.is_ipc_enabled

    # Test allocation/deallocation works
    buffer = mr.allocate(1024, stream=device.default_stream)
    assert buffer.size == 1024
    assert buffer.device_id == device.device_id
    buffer.close()


def test_pinned_memory_resource_initialization(init_cuda):
    device = Device()
    skip_if_pinned_memory_unsupported(device)

    device.set_current()

    mr = create_pinned_memory_resource_or_xfail(xfail_device=device)
    assert mr.is_device_accessible
    assert mr.is_host_accessible

    # Test allocation/deallocation works
    buffer = _allocate_pinned_buffer_or_xfail(mr, 1024, device=device)
    assert buffer.size == 1024
    assert buffer.device_id == -1  # Not bound to any GPU
    assert buffer.is_host_accessible
    assert buffer.memory_resource == mr
    assert buffer.is_device_accessible
    buffer.close()


@pytest.mark.agent_authored(model="cursor-grok-4.5")
def test_pinned_memory_resource_rejects_unsupported_host_pool(init_cuda):
    """allocate() must fail on devices without host memory pool support (see #2486)."""
    device = init_cuda
    if device.properties.host_memory_pools_supported:
        pytest.skip("Device supports host memory pools")

    try:
        mr = PinnedMemoryResource(PinnedMemoryResourceOptions(max_size=POOL_SIZE))
    except CUDAError as exc:
        if "CUDA_ERROR_NOT_SUPPORTED" in str(exc):
            pytest.skip("PinnedMemoryResource is not supported on this platform/device")
        raise
    try:
        with pytest.raises(RuntimeError, match="does not support.*LegacyPinnedMemoryResource"):
            mr.allocate(1024, stream=device.default_stream)
    finally:
        mr.close()


def test_managed_memory_resource_initialization(init_cuda):
    device = Device()
    skip_if_managed_memory_unsupported(device)

    device.set_current()

    mr = create_managed_memory_resource_or_skip()
    assert mr.is_device_accessible
    assert mr.is_host_accessible

    # Test allocation/deallocation works
    buffer = mr.allocate(1024, stream=device.default_stream)
    assert buffer.size == 1024
    assert buffer.is_host_accessible  # But accessible from host
    assert buffer.memory_resource == mr
    assert buffer.is_device_accessible
    buffer.close()


def get_handle_type():
    if IS_WINDOWS:
        return (("win32", None), ("win32_kmt", None))
    else:
        return (("posix_fd", None),)


@pytest.mark.parametrize("use_device_object", [True, False])
@pytest.mark.parametrize("handle_type", get_handle_type())
def test_vmm_allocator_basic_allocation(use_device_object, handle_type):
    """Test basic VMM allocation functionality.

    This test verifies that VirtualMemoryResource can allocate memory
    using CUDA VMM APIs with default configuration.
    """
    device = Device()
    device.set_current()

    # Skip if virtual memory management is not supported
    if not device.properties.virtual_memory_management_supported:
        pytest.skip("Virtual memory management is not supported on this device")

    handle_type, security_attribute = handle_type  # unpack
    options = VirtualMemoryResourceOptions(handle_type=handle_type)
    # Create VMM allocator with default config
    device_arg = device if use_device_object else device.device_id
    vmm_mr = VirtualMemoryResource(device_arg, config=options)

    # Test basic allocation
    try:
        buffer = vmm_mr.allocate(4096)
    except NotImplementedError:
        assert handle_type == "win32"
        return
    assert buffer.size >= 4096  # May be aligned up
    assert buffer.device_id == device.device_id
    assert buffer.memory_resource == vmm_mr

    # Test deallocation
    buffer.close()

    # Test multiple allocations
    buffers = []
    for i in range(5):
        buf = vmm_mr.allocate(1024 * (i + 1))
        buffers.append(buf)
        assert buf.size >= 1024 * (i + 1)

    # Clean up
    for buf in buffers:
        buf.close()


def test_vmm_allocator_policy_configuration():
    """Test VMM allocator with different policy configurations.

    This test verifies that VirtualMemoryResource can be configured
    with different allocation policies and that the configuration affects
    the allocation behavior.
    """
    device = Device()
    device.set_current()

    # Skip if virtual memory management is not supported
    if not device.properties.virtual_memory_management_supported:
        pytest.skip("Virtual memory management is not supported on this device")

    # Skip if GPU Direct RDMA is not supported
    if not device.properties.gpu_direct_rdma_supported:
        pytest.skip("This test requires a device that supports GPU Direct RDMA")

    # Test with custom VMM config
    custom_config = VirtualMemoryResourceOptions(
        allocation_type=VirtualMemoryAllocationType.PINNED,
        location_type=VirtualMemoryLocationType.DEVICE,
        granularity=VirtualMemoryGranularityType.MINIMUM,
        gpu_direct_rdma=True,
        handle_type=VirtualMemoryHandleType.POSIX_FD if not IS_WINDOWS else VirtualMemoryHandleType.WIN32_KMT,
        peers=(),
        self_access=VirtualMemoryAccessType.READ_WRITE,
        peer_access=VirtualMemoryAccessType.READ_WRITE,
    )

    vmm_mr = VirtualMemoryResource(device, config=custom_config)

    # Verify configuration is applied
    assert vmm_mr.config == custom_config
    assert vmm_mr.config.gpu_direct_rdma is True
    assert vmm_mr.config.granularity == "minimum"

    # Test allocation with custom config
    try:
        buffer = vmm_mr.allocate(8192)
    except CUDAError as exc:
        msg = str(exc)
        if "CUDA_ERROR_INVALID_DEVICE" in msg:
            pytest.xfail("TODO(#1300): Failing on Jetson AGX Orin P3730")
        raise
    assert buffer.size >= 8192
    assert buffer.device_id == device.device_id

    # Test policy modification
    new_config = VirtualMemoryResourceOptions(
        allocation_type="pinned",
        location_type="device",
        granularity="recommended",
        gpu_direct_rdma=False,
        handle_type="posix_fd" if not IS_WINDOWS else "win32_kmt",
        peers=(),
        self_access="r",  # Read-only access
        peer_access="r",
    )

    # Modify allocation policy
    try:
        modified_buffer = vmm_mr.modify_allocation(buffer, 16384, config=new_config)
    except CUDAError as exc:
        msg = str(exc)
        if "CUDA_ERROR_UNKNOWN" in msg:
            pytest.xfail("TODO(#1300): Known to fail already with CTK 13.0 (Windows)")
        raise
    assert modified_buffer.size >= 16384
    assert vmm_mr.config == new_config
    assert vmm_mr.config.self_access == "r"

    # Clean up
    modified_buffer.close()


@pytest.mark.parametrize("handle_type", get_handle_type())
def test_vmm_allocator_grow_allocation(handle_type):
    """Test VMM allocator's ability to grow existing allocations.

    This test verifies that VirtualMemoryResource can grow existing
    allocations while preserving the base pointer when possible.
    """
    device = Device()
    device.set_current()

    # Skip if virtual memory management is not supported (we need it for VMM)
    if not device.properties.virtual_memory_management_supported:
        pytest.skip("Virtual memory management is not supported on this device")

    handle_type, security_attribute = handle_type  # unpack
    options = VirtualMemoryResourceOptions(handle_type=handle_type)
    vmm_mr = VirtualMemoryResource(device, config=options)

    # Create initial allocation
    try:
        buffer = vmm_mr.allocate(2 * 1024 * 1024)
    except NotImplementedError:
        assert handle_type == "win32"
        return
    original_size = buffer.size

    # Grow the allocation
    grown_buffer = vmm_mr.modify_allocation(buffer, 4 * 1024 * 1024)

    # Verify growth
    assert grown_buffer.size >= 4 * 1024 * 1024
    assert grown_buffer.size > original_size
    # Because of the slow path, the pointer may change
    # We cannot assert that the new pointer is the same,
    # but we can assert that a new pointer was assigned
    assert grown_buffer.handle is not None

    # Test growing to same size (should return original buffer)
    same_buffer = vmm_mr.modify_allocation(grown_buffer, 4 * 1024 * 1024)
    assert same_buffer.size == grown_buffer.size

    # Test growing to smaller size (should return original buffer)
    smaller_buffer = vmm_mr.modify_allocation(grown_buffer, 2 * 1024 * 1024)
    assert smaller_buffer.size == grown_buffer.size

    # Clean up
    grown_buffer.close()


def test_vmm_allocator_grow_allocation_fast_path(init_cuda, monkeypatch):
    """Exercise the VMM grow fast path with mocked driver calls.

    The real driver usually rejects the adjacent reservation that reaches this
    path, so the test supplies that precondition by construction and verifies
    the successful commit bookkeeping.
    """
    device = Device()
    if not device.properties.virtual_memory_management_supported:
        pytest.skip("Virtual memory management is not supported on this device")

    vmm_mr = VirtualMemoryResource(
        device,
        config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"),
    )

    # Build the prop the same shape modify_allocation does, so the helper's
    # cuMemCreate / _build_access_descriptors path sees production-like input.
    prop = driver.CUmemAllocationProp()
    prop.type = driver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED
    prop.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
    prop.location.id = device.device_id

    SUCCESS = driver.CUresult.CUDA_SUCCESS
    NEW_HANDLE = 0xBEEF
    calls = []

    def fake_create(size, p, flags):
        calls.append(("create", size))
        return (SUCCESS, NEW_HANDLE)

    def fake_map(ptr, size, offset, handle, flags):
        calls.append(("map", ptr, size, handle))
        return (SUCCESS,)

    def fake_set_access(ptr, size, descs, count):
        calls.append(("set_access", ptr, size, count))
        return (SUCCESS,)

    # Rollback-only entry points: registered as undo actions but, on a
    # successful commit, must never be invoked.
    def fake_unmap(ptr, size):
        calls.append(("unmap", ptr, size))
        return (SUCCESS,)

    def fake_release(handle):
        calls.append(("release", handle))
        return (SUCCESS,)

    def fake_addr_free(ptr, size):
        calls.append(("addr_free", ptr, size))
        return (SUCCESS,)

    monkeypatch.setattr(driver, "cuMemCreate", fake_create)
    monkeypatch.setattr(driver, "cuMemMap", fake_map)
    monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access)
    monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap)
    monkeypatch.setattr(driver, "cuMemRelease", fake_release)
    monkeypatch.setattr(driver, "cuMemAddressFree", fake_addr_free)

    # A real Buffer carries C++-owned handle state we must not fabricate; the
    # helper only reads and writes buf._size, so a light stand-in suffices.
    class FakeBuffer:
        def __init__(self, size):
            self._size = size

    original_size = 2 * 1024 * 1024
    aligned_additional = 2 * 1024 * 1024
    new_size = original_size + aligned_additional
    new_ptr = 0x10_0000  # stand-in VA for the (mocked) contiguous extension

    buf = FakeBuffer(original_size)
    result = vmm_mr._grow_allocation_fast_path(buf, new_size, prop, aligned_additional, new_ptr)

    # Fast-path contract: same buffer object, size updated in place.
    assert result is buf
    assert buf._size == new_size

    # Successful commit: create, map, set access, and no rollback calls.
    assert [c[0] for c in calls] == ["create", "map", "set_access"]
    assert ("create", aligned_additional) in calls
    assert ("map", new_ptr, aligned_additional, NEW_HANDLE) in calls
    assert ("set_access", new_ptr, aligned_additional, 1) in calls


def test_vmm_allocator_rdma_unsupported_exception():
    """Test that VirtualMemoryResource throws an exception when RDMA is requested but device doesn't support it.

    This test verifies that the VirtualMemoryResource constructor throws a RuntimeError
    when gpu_direct_rdma=True is requested but the device doesn't support virtual memory management.
    """
    device = Device()
    device.set_current()

    # Skip if virtual memory management is not supported (we need it for VMM)
    if not device.properties.virtual_memory_management_supported:
        pytest.skip("Virtual memory management is not supported on this device")

    # Skip if GPU Direct RDMA is supported (we want to test the unsupported case)
    if device.properties.gpu_direct_rdma_supported:
        pytest.skip("This test requires a device that doesn't support GPU Direct RDMA")

    # Test that requesting RDMA on an unsupported device throws an exception
    options = VirtualMemoryResourceOptions(gpu_direct_rdma=True)
    with pytest.raises(RuntimeError, match="GPU Direct RDMA is not supported on this device"):
        VirtualMemoryResource(device, config=options)


def test_device_memory_resource_with_options(init_cuda):
    device = Device()
    if not device.properties.memory_pools_supported:
        pytest.skip("Device does not support mempool operations")

    device.set_current()

    # Test basic pool creation
    options = DeviceMemoryResourceOptions(max_size=POOL_SIZE)
    mr = DeviceMemoryResource(device, options=options)
    assert mr.device_id == device.device_id
    assert mr.is_device_accessible
    assert not mr.is_host_accessible
    assert not mr.is_ipc_enabled

    # Test allocation and deallocation
    buffer1 = mr.allocate(1024, stream=device.default_stream)
    assert buffer1.handle != 0
    assert buffer1.size == 1024
    assert buffer1.memory_resource == mr
    buffer1.close()

    # Test multiple allocations
    buffer1 = mr.allocate(1024, stream=device.default_stream)
    buffer2 = mr.allocate(2048, stream=device.default_stream)
    assert buffer1.handle != buffer2.handle
    assert buffer1.size == 1024
    assert buffer2.size == 2048
    buffer1.close()
    buffer2.close()

    # Test stream-based allocation
    stream = device.create_stream()
    buffer = mr.allocate(1024, stream=stream)
    assert buffer.handle != 0
    buffer.close(stream)

    # Test memory copying between buffers from same pool
    src_buffer = mr.allocate(64, stream=device.default_stream)
    dst_buffer = mr.allocate(64, stream=device.default_stream)
    stream = device.create_stream()
    src_buffer.copy_to(dst_buffer, stream=stream)
    device.sync()
    dst_buffer.close()
    src_buffer.close()
    # TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
    device.sync()


def test_pinned_memory_resource_with_options(init_cuda):
    device = Device()
    skip_if_pinned_memory_unsupported(device)

    device.set_current()

    # Test basic pool creation
    options = PinnedMemoryResourceOptions(max_size=POOL_SIZE)
    mr = PinnedMemoryResource(options)
    assert mr.device_id == -1  # Not bound to any GPU
    assert mr.is_device_accessible
    assert mr.is_host_accessible
    assert not mr.is_ipc_enabled

    # Test allocation and deallocation
    buffer1 = mr.allocate(1024, stream=device.default_stream)
    assert buffer1.handle != 0
    assert buffer1.size == 1024
    assert buffer1.memory_resource == mr
    buffer1.close()

    # Test multiple allocations
    buffer1 = mr.allocate(1024, stream=device.default_stream)
    buffer2 = mr.allocate(2048, stream=device.default_stream)
    assert buffer1.handle != buffer2.handle
    assert buffer1.size == 1024
    assert buffer2.size == 2048
    buffer1.close()
    buffer2.close()

    # Test stream-based allocation
    stream = device.create_stream()
    buffer = mr.allocate(1024, stream=stream)
    assert buffer.handle != 0
    buffer.close(stream)

    # Test memory copying between buffers from same pool
    src_buffer = mr.allocate(64, stream=device.default_stream)
    dst_buffer = mr.allocate(64, stream=device.default_stream)
    stream = device.create_stream()
    src_buffer.copy_to(dst_buffer, stream=stream)
    device.sync()
    dst_buffer.close()
    src_buffer.close()
    # TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
    device.sync()


def test_managed_memory_resource_with_options(init_cuda):
    device = Device()
    skip_if_managed_memory_unsupported(device)

    device.set_current()

    # Test basic pool creation
    options = ManagedMemoryResourceOptions()
    mr = create_managed_memory_resource_or_skip(options)
    assert mr.is_device_accessible
    assert mr.is_host_accessible
    assert not mr.is_ipc_enabled

    # Test allocation and deallocation
    buffer1 = mr.allocate(1024, stream=device.default_stream)
    assert buffer1.handle != 0
    assert buffer1.size == 1024
    assert buffer1.memory_resource == mr
    buffer1.close()

    # Test multiple allocations
    buffer1 = mr.allocate(1024, stream=device.default_stream)
    buffer2 = mr.allocate(2048, stream=device.default_stream)
    assert buffer1.handle != buffer2.handle
    assert buffer1.size == 1024
    assert buffer2.size == 2048
    buffer1.close()
    buffer2.close()

    # Test stream-based allocation
    stream = device.create_stream()
    buffer = mr.allocate(1024, stream=stream)
    assert buffer.handle != 0
    buffer.close(stream)

    # Test memory copying between buffers from same pool
    src_buffer = mr.allocate(64, stream=device.default_stream)
    dst_buffer = mr.allocate(64, stream=device.default_stream)
    stream = device.create_stream()
    src_buffer.copy_to(dst_buffer, stream=stream)
    device.sync()
    dst_buffer.close()
    src_buffer.close()


def test_managed_memory_resource_preferred_location_default(init_cuda):
    """preferred_location property returns None when no preference is set."""
    device = Device()
    skip_if_managed_memory_unsupported(device)
    device.set_current()

    mr = create_managed_memory_resource_or_skip()
    assert mr.preferred_location is None


def test_managed_memory_resource_preferred_location_device(init_cuda):
    """preferred_location returns ("device", ordinal) for device preference."""
    device = Device()
    skip_if_managed_memory_unsupported(device)
    device.set_current()

    # Legacy style
    opts = ManagedMemoryResourceOptions(preferred_location=device.device_id)
    mr = create_managed_memory_resource_or_skip(opts)
    assert mr.preferred_location == ("device", device.device_id)

    # Explicit style
    opts = ManagedMemoryResourceOptions(
        preferred_location=device.device_id,
        preferred_location_type=ManagedMemoryLocationType.DEVICE,
    )
    mr = create_managed_memory_resource_or_skip(opts)
    assert mr.preferred_location == ("device", device.device_id)


def test_managed_memory_resource_preferred_location_host(init_cuda):
    """preferred_location returns ("host", None) for host preference."""
    device = Device()
    skip_if_managed_memory_unsupported(device)
    device.set_current()

    # Legacy style
    opts = ManagedMemoryResourceOptions(preferred_location=-1)
    mr = create_managed_memory_resource_or_skip(opts)
    assert mr.preferred_location == ("host", None)

    # Explicit style
    opts = ManagedMemoryResourceOptions(preferred_location_type="host")
    mr = create_managed_memory_resource_or_skip(opts)
    assert mr.preferred_location == ("host", None)


def test_managed_memory_resource_preferred_location_host_numa(init_cuda):
    """preferred_location returns ("host_numa", id) for NUMA preference."""
    device = Device()
    skip_if_managed_memory_unsupported(device)
    device.set_current()

    numa_id = device.properties.host_numa_id
    if numa_id < 0:
        pytest.skip("System does not support NUMA")

    # Auto-resolved from current device
    opts = ManagedMemoryResourceOptions(preferred_location_type="host_numa")
    mr = create_managed_memory_resource_or_skip(opts)
    assert mr.preferred_location == ("host_numa", numa_id)

    # Explicit NUMA node ID
    opts = ManagedMemoryResourceOptions(
        preferred_location=numa_id,
        preferred_location_type="host_numa",
    )
    mr = create_managed_memory_resource_or_skip(opts)
    assert mr.preferred_location == ("host_numa", numa_id)


def test_managed_memory_resource_preferred_location_validation(init_cuda):
    """Invalid preferred_location combinations raise errors."""
    device = Device()
    skip_if_managed_memory_unsupported(device)
    device.set_current()

    # Invalid preferred_location_type
    with pytest.raises(ValueError, match="'invalid' is not a valid ManagedMemoryLocationType. Must be "):
        ManagedMemoryResource(
            ManagedMemoryResourceOptions(
                preferred_location_type="invalid",
            )
        )

    # "device" requires a non-negative int
    with pytest.raises(ValueError, match="must be a device ordinal"):
        ManagedMemoryResource(
            ManagedMemoryResourceOptions(
                preferred_location_type="device",
            )
        )
    with pytest.raises(ValueError, match="must be a device ordinal"):
        ManagedMemoryResource(
            ManagedMemoryResourceOptions(
                preferred_location=-1,
                preferred_location_type="device",
            )
        )

    # "host" requires preferred_location=None
    with pytest.raises(ValueError, match="must be None"):
        ManagedMemoryResource(
            ManagedMemoryResourceOptions(
                preferred_location=0,
                preferred_location_type="host",
            )
        )

    # "host_numa" rejects negative IDs
    with pytest.raises(ValueError, match="must be a NUMA node ID"):
        ManagedMemoryResource(
            ManagedMemoryResourceOptions(
                preferred_location=-1,
                preferred_location_type="host_numa",
            )
        )

    # Legacy mode rejects invalid negative values
    with pytest.raises(ValueError, match="preferred_location must be"):
        ManagedMemoryResource(
            ManagedMemoryResourceOptions(
                preferred_location=-2,
            )
        )


@pytest.mark.thread_unsafe(reason="Uses mock.")
def test_managed_memory_resource_host_numa_auto_resolve_failure(init_cuda):
    """host_numa with None raises RuntimeError when NUMA ID cannot be determined."""
    from unittest.mock import MagicMock, patch

    device = Device()
    skip_if_managed_memory_unsupported(device)
    device.set_current()

    mock_dev = MagicMock()
    mock_dev.properties.host_numa_id = -1

    with (
        patch("cuda.core._device.Device", return_value=mock_dev),
        pytest.raises(RuntimeError, match="Cannot determine host NUMA ID"),
    ):
        ManagedMemoryResource(
            ManagedMemoryResourceOptions(
                preferred_location_type="host_numa",
            )
        )


def test_mempool_ipc_errors(mempool_device):
    """Test error cases when IPC operations are disabled."""
    device = mempool_device
    options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=False)
    mr = DeviceMemoryResource(device, options=options)
    buffer = mr.allocate(64, stream=device.default_stream)
    ipc_error_msg = "Memory resource is not IPC-enabled"

    with pytest.raises(RuntimeError, match=ipc_error_msg):
        _ = mr.allocation_handle

    with pytest.raises(RuntimeError, match=ipc_error_msg):
        _ = buffer.ipc_descriptor

    with pytest.raises(RuntimeError, match=ipc_error_msg):
        handle = IPCBufferDescriptor._init(b"", 0)
        Buffer.from_ipc_descriptor(mr, handle, stream=device.default_stream)

    buffer.close()
    # TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
    device.sync()


def test_pinned_mempool_ipc_basic():
    """Test basic IPC functionality for PinnedMemoryResource."""
    device = Device()
    device.set_current()

    skip_if_pinned_memory_unsupported(device)

    if platform.system() == "Windows":
        pytest.skip("IPC not implemented for Windows")

    if not supports_ipc_mempool(device):
        pytest.skip("Driver rejects IPC-enabled mempool creation on this platform")

    # Test IPC-enabled PinnedMemoryResource creation
    options = PinnedMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
    mr = PinnedMemoryResource(options)
    assert mr.is_ipc_enabled
    assert mr.is_device_accessible
    assert mr.is_host_accessible
    assert mr.device_id == -1  # pinned memory is not device-specific
    assert mr.numa_id >= 0  # IPC requires a concrete NUMA node

    # Test allocation handle export
    alloc_handle = mr.allocation_handle
    assert alloc_handle is not None

    # Test buffer allocation
    buffer = mr.allocate(1024, stream=device.default_stream)
    assert buffer.size == 1024
    assert buffer.is_device_accessible
    assert buffer.is_host_accessible

    # Test IPC descriptor
    ipc_desc = buffer.ipc_descriptor
    assert ipc_desc is not None
    assert ipc_desc.size == 1024

    buffer.close()
    # TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
    device.sync()
    mr.close()


def test_pinned_mempool_ipc_errors():
    """Test error cases when IPC operations are disabled for PinnedMemoryResource."""
    device = Device()
    device.set_current()

    skip_if_pinned_memory_unsupported(device)

    # Test with IPC disabled (default)
    options = PinnedMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=False)
    mr = PinnedMemoryResource(options)
    assert not mr.is_ipc_enabled
    assert mr.device_id == -1
    assert mr.numa_id == -1  # Non-IPC uses OS-managed placement

    buffer = mr.allocate(64, stream=device.default_stream)
    ipc_error_msg = "Memory resource is not IPC-enabled"

    with pytest.raises(RuntimeError, match=ipc_error_msg):
        _ = mr.allocation_handle

    with pytest.raises(RuntimeError, match=ipc_error_msg):
        _ = buffer.ipc_descriptor

    with pytest.raises(RuntimeError, match=ipc_error_msg):
        handle = IPCBufferDescriptor._init(b"", 0)
        Buffer.from_ipc_descriptor(mr, handle, stream=device.default_stream)

    buffer.close()
    # TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
    device.sync()
    mr.close()


def test_pinned_mr_numa_id_default_no_ipc(init_cuda):
    """numa_id defaults to -1 (OS-managed) when IPC is disabled."""
    device = Device()
    skip_if_pinned_memory_unsupported(device)

    mr = create_pinned_memory_resource_or_xfail(PinnedMemoryResourceOptions(max_size=POOL_SIZE), xfail_device=device)
    assert mr.numa_id == -1
    mr.close()

    mr = create_pinned_memory_resource_or_xfail(
        PinnedMemoryResourceOptions(ipc_enabled=False, max_size=POOL_SIZE), xfail_device=device
    )
    assert mr.numa_id == -1
    mr.close()


def test_pinned_mr_numa_id_default_with_ipc(init_cuda):
    """numa_id is derived from the current device when IPC is enabled."""
    device = Device()
    skip_if_pinned_memory_unsupported(device)

    if platform.system() == "Windows":
        pytest.skip("IPC not implemented for Windows")
    if not supports_ipc_mempool(device):
        pytest.skip("Driver rejects IPC-enabled mempool creation on this platform")

    expected_numa_id = device.properties.host_numa_id
    if expected_numa_id < 0:
        pytest.skip("System does not support NUMA")

    mr = create_pinned_memory_resource_or_xfail(
        PinnedMemoryResourceOptions(ipc_enabled=True, max_size=POOL_SIZE), xfail_device=device
    )
    assert mr.numa_id == expected_numa_id
    mr.close()


def test_pinned_mr_numa_id_explicit(init_cuda):
    """Explicit numa_id is used regardless of ipc_enabled."""
    device = Device()
    skip_if_pinned_memory_unsupported(device)

    host_numa_id = device.properties.host_numa_id
    if host_numa_id < 0:
        pytest.skip("System does not support NUMA")

    mr = create_pinned_memory_resource_or_xfail(
        PinnedMemoryResourceOptions(numa_id=host_numa_id, max_size=POOL_SIZE), xfail_device=device
    )
    assert mr.numa_id == host_numa_id
    mr.close()

    if platform.system() == "Windows":
        pytest.skip("IPC not implemented for Windows")
    if not supports_ipc_mempool(device):
        pytest.skip("Driver rejects IPC-enabled mempool creation on this platform")

    mr = create_pinned_memory_resource_or_xfail(
        PinnedMemoryResourceOptions(ipc_enabled=True, numa_id=host_numa_id, max_size=POOL_SIZE),
        xfail_device=device,
    )
    assert mr.numa_id == host_numa_id
    mr.close()


def test_pinned_mr_numa_id_negative_error(init_cuda):
    """Negative numa_id raises ValueError."""
    device = Device()
    skip_if_pinned_memory_unsupported(device)

    with pytest.raises(ValueError, match="numa_id must be >= 0"):
        # uncapped-pool-ok: numa_id is validated before the pool is created
        PinnedMemoryResource(PinnedMemoryResourceOptions(numa_id=-1))

    with pytest.raises(ValueError, match="numa_id must be >= 0"):
        # uncapped-pool-ok: numa_id is validated before the pool is created
        PinnedMemoryResource(PinnedMemoryResourceOptions(numa_id=-42))


@pytest.mark.parametrize("ipc_enabled", [True, False])
@pytest.mark.parametrize(
    "property_name,expected_type",
    [
        ("reuse_follow_event_dependencies", bool),
        ("reuse_allow_opportunistic", bool),
        ("reuse_allow_internal_dependencies", bool),
        ("release_threshold", int),
        ("reserved_mem_current", int),
        ("reserved_mem_high", int),
        ("used_mem_current", int),
        ("used_mem_high", int),
    ],
)
def test_mempool_attributes(ipc_enabled, memory_resource_factory, property_name, expected_type):
    """Test all properties of memory pool attributes for all memory resource types."""
    MR, MRops = memory_resource_factory
    device = Device()

    if MR is DeviceMemoryResource and not device.properties.memory_pools_supported:
        pytest.skip("Device does not support mempool operations")
    elif MR is PinnedMemoryResource:
        skip_if_pinned_memory_unsupported(device)
    elif MR is ManagedMemoryResource:
        skip_if_managed_memory_unsupported(device)

    # ManagedMemoryResource does not support IPC
    if MR is ManagedMemoryResource and ipc_enabled:
        pytest.skip(f"{MR.__name__} does not support IPC")

    device.set_current()

    if platform.system() == "Windows":
        return  # IPC not implemented for Windows

    if ipc_enabled and not supports_ipc_mempool(device):
        pytest.skip("Driver rejects IPC-enabled mempool creation on this platform")

    if MR is DeviceMemoryResource:
        options = MRops(max_size=POOL_SIZE, ipc_enabled=ipc_enabled)
        mr = MR(device, options=options)
        assert mr.is_ipc_enabled == ipc_enabled
    elif MR is PinnedMemoryResource:
        options = MRops(max_size=POOL_SIZE, ipc_enabled=ipc_enabled)
        mr = MR(options)
        assert mr.is_ipc_enabled == ipc_enabled
    elif MR is ManagedMemoryResource:
        options = MRops()
        mr = create_managed_memory_resource_or_skip(options)
        assert not mr.is_ipc_enabled

    # Get the property value
    value = getattr(mr.attributes, property_name)

    # Test type
    assert isinstance(value, expected_type), f"{property_name} should return {expected_type}, got {type(value)}"

    # Test value constraints
    if expected_type is int:
        assert value >= 0, f"{property_name} should be non-negative"

    # Test memory usage properties with actual allocations
    if property_name in ["reserved_mem_current", "used_mem_current"]:
        # Allocate some memory and check if values increase
        initial_value = value
        buffer = None
        try:
            buffer = mr.allocate(1024, stream=device.default_stream)
            new_value = getattr(mr.attributes, property_name)
            assert new_value >= initial_value, f"{property_name} should increase or stay same after allocation"
        finally:
            if buffer is not None:
                buffer.close()

    # Test high watermark properties
    if property_name in ["reserved_mem_high", "used_mem_high"]:
        # High watermark should never be less than current
        current_prop = property_name.replace("_high", "_current")
        current_value = getattr(mr.attributes, current_prop)
        assert value >= current_value, f"{property_name} should be >= {current_prop}"


def test_mempool_attributes_repr(memory_resource_factory):
    """Test the repr of memory pool attributes for all memory resource types."""
    MR, MRops = memory_resource_factory
    device = Device()

    if MR is DeviceMemoryResource and not device.properties.memory_pools_supported:
        pytest.skip("Device does not support mempool operations")
    elif MR is PinnedMemoryResource:
        skip_if_pinned_memory_unsupported(device)
    elif MR is ManagedMemoryResource:
        skip_if_managed_memory_unsupported(device)

    device.set_current()

    if MR is DeviceMemoryResource:
        mr = MR(device, options={"max_size": 2048})
    elif MR is PinnedMemoryResource:
        mr = MR(options={"max_size": 2048})
    elif MR is ManagedMemoryResource:
        mr = create_managed_memory_resource_or_skip(options={})

    buffer1 = mr.allocate(64, stream=device.default_stream)
    buffer2 = mr.allocate(64, stream=device.default_stream)
    buffer1.close()
    assert re.match(
        r".*Attributes\(release_threshold=\d+, reserved_mem_current=\d+, reserved_mem_high=\d+, "
        r"reuse_allow_internal_dependencies=(True|False), reuse_allow_opportunistic=(True|False), "
        r"reuse_follow_event_dependencies=(True|False), used_mem_current=\d+, used_mem_high=\d+\)",
        str(mr.attributes),
    )
    buffer2.close()


def test_mempool_attributes_ownership(memory_resource_factory):
    """Ensure the attributes bundle keeps the pool alive via the handle."""
    MR, MRops = memory_resource_factory
    device = Device()

    if MR is DeviceMemoryResource and not device.properties.memory_pools_supported:
        pytest.skip("Device does not support mempool operations")
    elif MR is PinnedMemoryResource:
        skip_if_pinned_memory_unsupported(device)
    elif MR is ManagedMemoryResource:
        skip_if_managed_memory_unsupported(device)

    # Skip if IPC mempool is not supported on this platform/device (only relevant for DeviceMemoryResource)
    if MR is DeviceMemoryResource and not supports_ipc_mempool(device):
        pytest.skip("Driver rejects IPC-enabled mempool creation on this platform")

    device.set_current()

    if MR is DeviceMemoryResource:
        mr = MR(device, {"max_size": POOL_SIZE})
    elif MR is PinnedMemoryResource:
        mr = MR({"max_size": POOL_SIZE})
    elif MR is ManagedMemoryResource:
        mr = create_managed_memory_resource_or_skip({})

    attributes = mr.attributes
    mr.close()
    del mr

    # The attributes bundle keeps the pool alive via MemoryPoolHandle,
    # so accessing attributes still works even after the MR is deleted.
    _ = attributes.used_mem_high  # Should not raise


# Ensure that memory views dellocate their reference to dlpack tensors
@pytest.mark.skipif(np is None, reason="numpy is not installed")
def test_strided_memory_view_leak():
    arr = np.zeros(1048576, dtype=np.uint8)
    before = sys.getrefcount(arr)
    for idx in range(10):
        StridedMemoryView.from_any_interface(arr, stream_ptr=-1)
    after = sys.getrefcount(arr)
    assert before == after


def test_strided_memory_view_refcnt():
    # Use Fortran ordering so strides is used
    a = np.zeros((64, 4), dtype=np.uint8, order="F")
    av = StridedMemoryView.from_any_interface(a, stream_ptr=-1)
    # segfaults if refcnt is wrong
    assert av.shape[0] == 64
    assert sys.getrefcount(av.shape) >= 2

    assert av.strides[0] == 1
    assert av.strides[1] == 64
    assert sys.getrefcount(av.strides) >= 2


def test_graph_memory_resource_object(init_cuda):
    device = Device()
    gmr1 = GraphMemoryResource(device)
    gmr2 = GraphMemoryResource(device)
    gmr3 = GraphMemoryResource(device.device_id)

    # These objects are interned.
    assert gmr1 is gmr2 is gmr3
    assert gmr1 == gmr2 == gmr3


@pytest.mark.skipif(np is None, reason="numpy is not installed")
def test_strided_memory_view_dlpack_errors():
    arr = np.zeros(64, dtype=np.uint8)
    smv = StridedMemoryView.from_any_interface(arr, stream_ptr=-1)
    with pytest.raises(BufferError, match="dl_device other than None"):
        smv.__dlpack__(dl_device=())
    with pytest.raises(BufferError, match="copy=True"):
        smv.__dlpack__(copy=True)
    with pytest.raises(BufferError, match="Expected max_version"):
        smv.__dlpack__(max_version=(9, 8, 7))


def test_memory_resource_alloc_zero_bytes(init_cuda, memory_resource_factory):
    MR, MROps = memory_resource_factory

    device = Device()
    device.set_current()

    if MR is DeviceMemoryResource and not device.properties.memory_pools_supported:
        pytest.skip("Device does not support mempool operations")
    elif MR is PinnedMemoryResource:
        skip_if_pinned_memory_unsupported(device)
        mr = create_pinned_memory_resource_or_xfail(xfail_device=device)
    elif MR is ManagedMemoryResource:
        skip_if_managed_memory_unsupported(device)
        mr = create_managed_memory_resource_or_skip(MROps(preferred_location=device.device_id))
    else:
        assert MR is DeviceMemoryResource
        mr = MR(device)

    buffer = mr.allocate(0, stream=device.default_stream)
    device.sync()
    assert buffer.handle >= 0
    assert buffer.size == 0
    assert buffer.device_id == mr.device_id


def test_strided_memory_view_not_in_top_level_namespace():
    # Regression test for issue #2027: StridedMemoryView and
    # args_viewable_as_strided_memory must only be exposed via
    # cuda.core.utils, not the top-level cuda.core namespace.
    import cuda.core
    import cuda.core.utils

    assert not hasattr(cuda.core, "StridedMemoryView")
    assert not hasattr(cuda.core, "args_viewable_as_strided_memory")

    assert hasattr(cuda.core.utils, "StridedMemoryView")
    assert hasattr(cuda.core.utils, "args_viewable_as_strided_memory")


def test_top_level_namespace_excludes_known_leaks():
    # Hardening test: lock the public top-level namespace against
    # accidental re-introduction of known-leaked symbols.
    import cuda.core

    public = {n for n in dir(cuda.core) if not n.startswith("_")}
    leaked = {"StridedMemoryView", "args_viewable_as_strided_memory"}
    assert not (public & leaked)


def test_legacy_pinned_allocate_zero_size(init_cuda):
    """LegacyPinnedMemoryResource.allocate(0) skips the driver call and uses ptr=0."""
    mr = LegacyPinnedMemoryResource()
    buf = mr.allocate(0)
    assert buf.size == 0
    # No driver call was made; handle is the sentinel 0.
    assert int(buf.handle) == 0


def test_legacy_pinned_device_id_raises():
    """LegacyPinnedMemoryResource.device_id raises; pinned memory is not bound to a GPU."""
    mr = LegacyPinnedMemoryResource()
    with pytest.raises(RuntimeError, match="not bound to any GPU"):
        _ = mr.device_id


def test_synchronous_memory_resource_basic(init_cuda):
    """_SynchronousMemoryResource exercises properties and allocate paths (zero, non-zero, with-stream)."""
    from cuda.core._memory._legacy import _SynchronousMemoryResource

    dev = Device()
    mr = _SynchronousMemoryResource(dev.device_id)
    assert mr.is_device_accessible is True
    assert mr.is_host_accessible is False
    assert mr.device_id == dev.device_id

    # Zero-size allocation takes the ptr=0 fast path.
    zero_buf = mr.allocate(0)
    assert zero_buf.size == 0
    assert int(zero_buf.handle) == 0
    zero_buf.close(stream=None)

    # Non-zero allocation goes through cuMemAlloc; close with stream=None for
    # the simple path. The explicit-stream close path is covered separately.
    buf = mr.allocate(64)
    assert buf.size == 64
    assert int(buf.handle) != 0
    buf.close(stream=None)

    # allocate(size, stream=stream) exercises Stream_accept validation on the
    # allocate side (cuMemAlloc is synchronous so the stream is accepted but unused).
    stream = dev.create_stream()
    buf2 = mr.allocate(32, stream=stream)
    assert buf2.size == 32
    assert int(buf2.handle) != 0
    buf2.close(stream=None)
    stream.close()


def test_synchronous_memory_resource_deallocate_accepts_stream(init_cuda):
    """_SynchronousMemoryResource.deallocate accepts an explicit stream."""
    from cuda.core._memory._legacy import _SynchronousMemoryResource

    dev = Device()
    mr = _SynchronousMemoryResource(dev.device_id)
    buf = mr.allocate(64)
    stream = dev.create_stream()
    buf.close(stream=stream)
    stream.close()


@pytest.mark.parametrize(
    ("method", "spec", "match"),
    [
        ("_access_to_flags", "bogus", "Unknown access spec"),
        ("_allocation_type_to_driver", "bogus", "Unsupported allocation_type"),
        ("_location_type_to_driver", "bogus", "Unsupported location_type"),
        ("_handle_type_to_driver", "bogus", "Unsupported handle_type"),
        ("_granularity_to_driver", "bogus", "Unsupported granularity"),
    ],
)
def test_vmm_options_spec_validators_raise(method, spec, match):
    """Every VMM spec validator static method rejects unknown strings with ValueError."""
    fn = getattr(VirtualMemoryResourceOptions, method)
    with pytest.raises(ValueError, match=match):
        fn(spec)


def test_vmm_options_handle_type_win32_raises():
    """_handle_type_to_driver raises NotImplementedError for 'win32'."""
    with pytest.raises(NotImplementedError, match="win32 is currently not supported"):
        VirtualMemoryResourceOptions._handle_type_to_driver("win32")


@pytest.mark.agent_authored(model="claude-opus-5")
@pytest.mark.parametrize("location_type", ["host", "host_numa", "host_numa_current"])
def test_vmm_host_location_types_report_host_accessible(location_type):
    """Every host-backed location type reports is_host_accessible.

    __init__ classifies "host", "host_numa" and "host_numa_current" alike when
    deciding the resource is not bound to a device, so is_host_accessible must
    agree; otherwise a NUMA-located resource claims to be neither host- nor
    device-accessible.
    """
    device = Device()
    device.set_current()
    mr = VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type))
    assert mr.device is None
    assert mr.is_host_accessible is True


def test_device_memory_resource_peer_accessible_by_non_owned(mempool_device):
    """peer_accessible_by on a non-owned (default) DMR queries the driver live."""
    dev = mempool_device
    # The default DeviceMemoryResource(device) wraps the current device's
    # default pool, i.e. _mempool_owned is False, so accessing
    # peer_accessible_by exercises the live _DMR_query_peer_access path.
    mr = DeviceMemoryResource(dev)
    peers = mr.peer_accessible_by
    assert all(isinstance(p, Device) for p in peers)
    # __contains__ accepts int dev_ids; the owning device is never a peer.
    assert dev.device_id not in peers


def test_dmr_mempool_get_access_self(mempool_device):
    """DMR_mempool_get_access returns 'rw' when querying the owning device itself."""
    from cuda.core._memory._device_memory_resource import DMR_mempool_get_access

    mr = DeviceMemoryResource(mempool_device)
    # The owning device always has read-write access to its own pool.
    assert DMR_mempool_get_access(mr, mempool_device.device_id) == "rw"


def test_dmr_mempool_get_access_peer(mempool_device_x2):
    """DMR_mempool_get_access reflects peer access state for a different device."""
    from cuda.core._memory._device_memory_resource import DMR_mempool_get_access

    dev, peer = mempool_device_x2
    # Owned pool avoids peer-access contamination from other tests; max_size
    # caps VA to dodge Windows MCDM OOM
    mr = DeviceMemoryResource(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE))

    # Fresh owned pool: peer has no access.
    assert DMR_mempool_get_access(mr, peer.device_id) == ""
    # After granting access, peer has read-write.
    mr.peer_accessible_by = [peer]
    assert DMR_mempool_get_access(mr, peer.device_id) == "rw"
    # After revoking, peer is back to no access.
    mr.peer_accessible_by = []
    assert DMR_mempool_get_access(mr, peer.device_id) == ""


def test_dmr_peer_accessible_by_setter_empty(mempool_device):
    """Assigning an empty peer-access set to a fresh owned pool is a no-op."""
    # max_size caps VA to dodge Windows MCDM OOM
    mr = DeviceMemoryResource(mempool_device, options=DeviceMemoryResourceOptions(max_size=POOL_SIZE))
    assert set(mr.peer_accessible_by) == set()
    mr.peer_accessible_by = []
    assert set(mr.peer_accessible_by) == set()


@pytest.mark.agent_authored(model="claude-opus-4.8")
def test_mempool_attributes_cannot_instantiate_directly():
    """_MemPoolAttributes cannot be instantiated directly."""
    from cuda.core._memory._memory_pool import _MemPoolAttributes

    with pytest.raises(RuntimeError, match="cannot be instantiated directly"):
        _MemPoolAttributes()


@pytest.mark.agent_authored(model="claude-opus-4.8")
def test_dmr_handle_and_ownership(mempool_device):
    """An options-created pool is handle-owning with a live handle; wrapping the device's current pool is non-owning."""
    owned = DeviceMemoryResource(mempool_device, DeviceMemoryResourceOptions(max_size=POOL_SIZE))
    assert owned.is_handle_owned is True
    handle = owned.handle
    assert handle is not None
    assert int(handle) != 0

    non_owned = DeviceMemoryResource(mempool_device)
    assert non_owned.is_handle_owned is False


@pytest.mark.agent_authored(model="claude-opus-4.8")
def test_dmr_deallocate_frees_pool_pointer(mempool_device):
    """Closing a Buffer.from_handle(..., mr=mr) view frees the pointer via the Python
    _MemPool.deallocate path; the pool's in-use bytes drop back."""
    dev = mempool_device
    stream = dev.default_stream
    mr = DeviceMemoryResource(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE))
    size = 256
    # Raw pool allocation owned by nobody else, so exactly one owner frees it (no
    # double free); a Buffer.from_handle view then routes teardown through the
    # Python deallocate path that mr.allocate()'s C++-direct free would skip.
    ptr = handle_return(driver.cuMemAllocFromPoolAsync(size, mr.handle, stream.handle))
    stream.sync()
    used_after_alloc = mr.attributes.used_mem_current
    assert used_after_alloc >= size
    buf = Buffer.from_handle(int(ptr), size, mr=mr)
    buf.close(stream)
    stream.sync()
    assert int(buf.handle) == 0
    # In-use bytes fell back, so the pointer was actually returned (buf.handle == 0
    # alone wouldn't prove it: the deleter callback swallows a failed free).
    assert mr.attributes.used_mem_current < used_after_alloc


@pytest.mark.agent_authored(model="claude-opus-4.8")
def test_dmr_close_is_idempotent(mempool_device):
    """Closing an owned DeviceMemoryResource twice is safe (the second close is a no-op)."""
    mr = DeviceMemoryResource(mempool_device, DeviceMemoryResourceOptions(max_size=POOL_SIZE))
    assert mr.is_handle_owned is True
    assert int(mr.handle) != 0
    mr.close()
    # First close releases the pool handle itself, not just ownership.
    assert int(mr.handle) == 0
    assert mr.is_handle_owned is False
    mr.close()  # no-op on the now-null handle
    assert int(mr.handle) == 0
    assert mr.is_handle_owned is False


@pytest.mark.agent_authored(model="claude-opus-4.8")
def test_dmr_ipc_enabled_unsupported_raises(mempool_device):
    """Requesting an IPC-enabled pool where memory IPC is unsupported raises RuntimeError."""
    if not IS_WINDOWS:
        pytest.skip("memory IPC is supported on this platform; unsupported-raise path is Windows-only")
    with pytest.raises(RuntimeError, match="IPC is not available"):
        # uncapped-pool-ok: IPC support is checked before the pool is created
        DeviceMemoryResource(mempool_device, DeviceMemoryResourceOptions(ipc_enabled=True))
