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

import ctypes

import helpers
from cuda_python_test_helpers.marks import requires_module
from helpers.misc import StreamWrapper

try:
    import cupy as cp
except ImportError:
    cp = None
import numpy as np
import pytest
from conftest import skipif_need_cuda_headers

from cuda.core import (
    Device,
    DeviceMemoryResource,
    LaunchConfig,
    LegacyPinnedMemoryResource,
    Program,
    ProgramOptions,
    launch,
)
from cuda.core._memory._legacy import _SynchronousMemoryResource
from cuda.core._utils.cuda_utils import CUDAError
from cuda.core.typing import ObjectCodeFormatType, SourceCodeType


def test_launch_config_init(init_cuda):
    config = LaunchConfig(grid=(1, 1, 1), block=(1, 1, 1), shmem_size=0)
    assert config.grid == (1, 1, 1)
    assert config.block == (1, 1, 1)
    assert config.shmem_size == 0

    config = LaunchConfig(grid=(2, 2, 2), block=(2, 2, 2), shmem_size=1024)
    assert config.grid == (2, 2, 2)
    assert config.block == (2, 2, 2)
    assert config.shmem_size == 1024


def test_launch_config_invalid_values():
    with pytest.raises(ValueError):
        LaunchConfig(grid=0, block=1)

    with pytest.raises(ValueError):
        LaunchConfig(grid=(0, 1), block=1)

    with pytest.raises(ValueError):
        LaunchConfig(grid=(1, 1, 1), block=0)

    with pytest.raises(ValueError):
        LaunchConfig(grid=(1, 1, 1), block=(0, 1))


def test_launch_config_shmem_size():
    config = LaunchConfig(grid=(1, 1, 1), block=(1, 1, 1), shmem_size=2048)
    assert config.shmem_size == 2048

    config = LaunchConfig(grid=(1, 1, 1), block=(1, 1, 1))
    assert config.shmem_size == 0


def test_launch_config_cluster_grid_conversion(init_cuda):
    """Test that LaunchConfig preserves original grid values and conversion happens in native config."""
    try:
        # Test case 1: 1D - Issue #867 example
        config = LaunchConfig(grid=4, cluster=2, block=32)
        assert config.grid == (4, 1, 1), f"Expected (4, 1, 1), got {config.grid}"
        assert config.cluster == (2, 1, 1), f"Expected (2, 1, 1), got {config.cluster}"
        assert config.block == (32, 1, 1), f"Expected (32, 1, 1), got {config.block}"

        # Test case 2: 2D grid and cluster
        config = LaunchConfig(grid=(2, 3), cluster=(2, 2), block=32)
        assert config.grid == (2, 3, 1), f"Expected (2, 3, 1), got {config.grid}"
        assert config.cluster == (2, 2, 1), f"Expected (2, 2, 1), got {config.cluster}"

        # Test case 3: 3D full specification
        config = LaunchConfig(grid=(2, 2, 2), cluster=(3, 3, 3), block=(8, 8, 8))
        assert config.grid == (2, 2, 2), f"Expected (2, 2, 2), got {config.grid}"
        assert config.cluster == (3, 3, 3), f"Expected (3, 3, 3), got {config.cluster}"

        # Test case 4: Identity case
        config = LaunchConfig(grid=1, cluster=1, block=32)
        assert config.grid == (1, 1, 1), f"Expected (1, 1, 1), got {config.grid}"

        # Test case 5: No cluster (should not convert grid)
        config = LaunchConfig(grid=4, block=32)
        assert config.grid == (4, 1, 1), f"Expected (4, 1, 1), got {config.grid}"
        assert config.cluster is None

    except CUDAError:
        pytest.skip("Driver or GPU not new enough for thread block clusters")


def test_launch_config_native_conversion(init_cuda):
    """Test that _to_native_launch_config correctly converts grid from cluster units to block units."""
    from cuda.core._launch_config import _to_native_launch_config

    try:
        # Test case 1: 1D - Issue #867 example
        config = LaunchConfig(grid=4, cluster=2, block=32)
        native_config = _to_native_launch_config(config)
        assert native_config.gridDimX == 8, f"Expected gridDimX=8, got {native_config.gridDimX}"
        assert native_config.gridDimY == 1, f"Expected gridDimY=1, got {native_config.gridDimY}"
        assert native_config.gridDimZ == 1, f"Expected gridDimZ=1, got {native_config.gridDimZ}"

        # Test case 2: 2D grid and cluster
        config = LaunchConfig(grid=(2, 3), cluster=(2, 2), block=32)
        native_config = _to_native_launch_config(config)
        assert native_config.gridDimX == 4, f"Expected gridDimX=4, got {native_config.gridDimX}"
        assert native_config.gridDimY == 6, f"Expected gridDimY=6, got {native_config.gridDimY}"
        assert native_config.gridDimZ == 1, f"Expected gridDimZ=1, got {native_config.gridDimZ}"

        # Test case 3: No cluster (should not convert grid)
        config = LaunchConfig(grid=4, block=32)
        native_config = _to_native_launch_config(config)
        assert native_config.gridDimX == 4, f"Expected gridDimX=4, got {native_config.gridDimX}"
        assert native_config.gridDimY == 1, f"Expected gridDimY=1, got {native_config.gridDimY}"
        assert native_config.gridDimZ == 1, f"Expected gridDimZ=1, got {native_config.gridDimZ}"

    except CUDAError:
        pytest.skip("Driver or GPU not new enough for thread block clusters")


def test_to_native_launch_config_no_cluster():
    """Covers the no-cluster path of _to_native_launch_config; no Hopper+ required."""
    from cuda.core._launch_config import _to_native_launch_config

    config = LaunchConfig(grid=(4, 5, 6), block=(7, 8, 9), shmem_size=128)
    native = _to_native_launch_config(config)
    assert native.gridDimX == 4, f"Expected gridDimX=4, got {native.gridDimX}"
    assert native.gridDimY == 5, f"Expected gridDimY=5, got {native.gridDimY}"
    assert native.gridDimZ == 6, f"Expected gridDimZ=6, got {native.gridDimZ}"
    assert native.blockDimX == 7, f"Expected blockDimX=7, got {native.blockDimX}"
    assert native.blockDimY == 8, f"Expected blockDimY=8, got {native.blockDimY}"
    assert native.blockDimZ == 9, f"Expected blockDimZ=9, got {native.blockDimZ}"
    assert native.sharedMemBytes == 128, f"Expected sharedMemBytes=128, got {native.sharedMemBytes}"
    assert native.numAttrs == 0, f"Expected numAttrs=0, got {native.numAttrs}"
    assert list(native.attrs) == [], f"Expected empty attrs, got {list(native.attrs)}"


def test_launch_config_cooperative_unsupported(monkeypatch):
    """LaunchConfig(is_cooperative=True) raises when device does not support it."""
    from cuda.core import _launch_config as _lc_mod

    class _FakeProps:
        cooperative_launch = False

    class _FakeDev:
        properties = _FakeProps()

    monkeypatch.setattr(_lc_mod, "Device", lambda: _FakeDev())
    with pytest.raises(CUDAError, match="cooperative kernels are not supported"):
        LaunchConfig(grid=1, block=1, is_cooperative=True)


def test_to_native_launch_config_cooperative(monkeypatch):
    """Covers the is_cooperative branch of _to_native_launch_config; Device is mocked so it runs on any GPU."""
    from cuda.bindings import driver
    from cuda.core import _launch_config as _lc_mod
    from cuda.core._launch_config import _to_native_launch_config

    class _FakeProps:
        cooperative_launch = True

    class _FakeDev:
        properties = _FakeProps()

    monkeypatch.setattr(_lc_mod, "Device", lambda: _FakeDev())

    config = LaunchConfig(grid=2, block=4, is_cooperative=True)
    native = _to_native_launch_config(config)
    assert native.gridDimX == 2
    assert native.blockDimX == 4
    assert native.numAttrs == 1
    attr = native.attrs[0]
    assert attr.id == driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_COOPERATIVE, (
        f"Expected CU_LAUNCH_ATTRIBUTE_COOPERATIVE, got {attr.id}"
    )
    assert attr.value.cooperative == 1, f"Expected cooperative=1, got {attr.value.cooperative}"


def test_to_native_launch_config_pdl():
    """LaunchConfig(programmatic_stream_serialization=True) maps to the PDL launch attribute."""
    from cuda.bindings import driver
    from cuda.core._launch_config import _to_native_launch_config

    config = LaunchConfig(grid=2, block=4, programmatic_stream_serialization=True)
    native = _to_native_launch_config(config)
    assert native.gridDimX == 2
    assert native.blockDimX == 4
    assert native.numAttrs == 1
    attr = native.attrs[0]
    assert attr.id == driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION, (
        f"Expected CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION, got {attr.id}"
    )
    assert attr.value.programmaticStreamSerializationAllowed == 1, (
        f"Expected programmaticStreamSerializationAllowed=1, got {attr.value.programmaticStreamSerializationAllowed}"
    )


@skipif_need_cuda_headers
def test_pdl_primary_secondary_overlap_same_stream():
    """Primary + secondary PDL launch on one stream can overlap on Hopper+.

    Secondary is launched with ``programmatic_stream_serialization=True``. After
    the primary triggers completion, it spins until it observes a flag written by
    the secondary's independent preamble — proving both grids were resident at
    once. Without PDL, the secondary cannot start until the primary exits.

    Note concurrency is opportunistic, so a missing overlap execution is reported as
    an expected failure.
    """
    dev = Device()
    if dev.compute_capability < (9, 0):
        pytest.skip("Programmatic Dependent Launch requires compute capability >= 9.0")
    dev.set_current()
    stream = dev.create_stream(options={"nonblocking": True})

    # clock64 budgets are in GPU cycles; keep the post-trigger window long enough
    # for the secondary to boot, but short enough for a unit test.
    code = r"""
    #include <cuda_device_runtime_api.h>

    extern "C" __global__ void primary_kernel(int* secondary_started, int* overlapped) {
        cudaTriggerProgrammaticLaunchCompletion();

        const long long deadline = clock64() + 100000000LL;  // ~50ms @ ~2GHz
        if (threadIdx.x == 0 && blockIdx.x == 0) {
            while (clock64() < deadline) {
                if (atomicAdd(secondary_started, 0) != 0) {
                    atomicExch(overlapped, 1);
                    return;
                }
                __nanosleep(1000);
            }
        }
    }

    extern "C" __global__ void secondary_kernel(int* secondary_started) {
        if (threadIdx.x == 0 && blockIdx.x == 0) {
            atomicExch(secondary_started, 1);
        }
    }
    """

    arch = "".join(f"{i}" for i in dev.compute_capability)
    pro_opts = ProgramOptions(std="c++17", arch=f"sm_{arch}", include_path=helpers.CUDA_INCLUDE_PATH)
    prog = Program(code, code_type="c++", options=pro_opts)
    mod = prog.compile("cubin")
    primary = mod.get_kernel("primary_kernel")
    secondary = mod.get_kernel("secondary_kernel")

    mr = LegacyPinnedMemoryResource()
    secondary_started = np.from_dlpack(mr.allocate(4)).view(np.int32)
    overlapped = np.from_dlpack(mr.allocate(4)).view(np.int32)

    primary_cfg = LaunchConfig(grid=1, block=1)
    secondary_cfg = LaunchConfig(grid=1, block=1, programmatic_stream_serialization=True)
    secondary_serial_cfg = LaunchConfig(grid=1, block=1)

    def _run(secondary_launch_cfg: LaunchConfig) -> int:
        secondary_started[0] = 0
        overlapped[0] = 0
        launch(stream, primary_cfg, primary, secondary_started.ctypes.data, overlapped.ctypes.data)
        launch(stream, secondary_launch_cfg, secondary, secondary_started.ctypes.data)
        stream.sync()
        return int(overlapped[0])

    # Without the PDL attribute, same-stream kernels stay serialized.
    assert _run(secondary_serial_cfg) == 0, "Expected no overlap when programmatic_stream_serialization is False"

    # PDL overlap is opportunistic; retry a few times on a quiet GPU.
    saw_overlap = False
    for _ in range(5):
        if _run(secondary_cfg) == 1:
            saw_overlap = True
            break

    if not saw_overlap:
        # Overlap is never guaranteed by the driver, so a miss is reported as an
        # expected failure rather than turning a busy GPU into a red CI run.
        pytest.xfail(
            "PDL (Programmatic Dependent Launch) overlap was not observed. "
            "If this keeps xfailing in CI, manually re-check on a quiet Hopper+ GPU."
        )

    print(
        f"PDL (Programmatic Dependent Launch) overlap verified on {dev.name} compute capability {dev.compute_capability}",
        flush=True,
    )


def test_launch_config_cluster_accepts_hopper_cc(monkeypatch):
    """LaunchConfig accepts ``cluster`` when the device reports compute
    capability >= 9.0. Device is mocked so the cluster-cast branch runs on any
    GPU (real cluster support otherwise requires Hopper+)."""
    from cuda.core import _launch_config as _lc_mod

    class _FakeDev:
        compute_capability = (9, 0)

    # looked_up confirms the mock took effect.
    looked_up = []
    monkeypatch.setattr(_lc_mod, "Device", lambda: looked_up.append(1) or _FakeDev())

    config = LaunchConfig(grid=(2, 3), cluster=(2, 2), block=32)
    assert looked_up, "Device was not looked up via the module global; mock did not take effect"
    assert config.cluster == (2, 2, 1)
    assert config.grid == (2, 3, 1)


def test_launch_config_cluster_rejects_pre_hopper_cc(monkeypatch):
    """LaunchConfig(cluster=...) raises on a device with compute capability < 9.0."""
    from cuda.core import _launch_config as _lc_mod

    class _FakeDev:
        compute_capability = (8, 6)

    # looked_up confirms the mock took effect.
    looked_up = []
    monkeypatch.setattr(_lc_mod, "Device", lambda: looked_up.append(1) or _FakeDev())

    with pytest.raises(CUDAError, match="thread block clusters are not supported"):
        LaunchConfig(grid=2, cluster=2, block=32)
    assert looked_up, "Device was not looked up via the module global; mock did not take effect"


def test_to_native_launch_config_cluster_branch():
    """Covers the cluster branch of ``_to_native_launch_config`` (grid is
    converted from cluster units to block units, plus the cluster-dimension
    attribute) without requiring Hopper.

    The cc gate lives in ``LaunchConfig.__init__``; ``cluster`` itself is a
    public attribute, so setting it on a cluster-free config yields the exact
    object ``__init__`` would build on Hopper and lets the conversion run on
    any GPU.

    Note: this exercises the standalone ``cpdef _to_native_launch_config``
    function (a duplicate of the ``LaunchConfig._to_native_launch_config``
    cdef method, slated for removal once all modules are cythonized), not the
    cdef method that ``launch`` / ``Module`` actually call in production.
    """
    from cuda.bindings import driver
    from cuda.core._launch_config import _to_native_launch_config

    config = LaunchConfig(grid=(2, 3, 4), block=(5, 6, 7))
    config.cluster = (2, 2, 2)
    native = _to_native_launch_config(config)

    # grid (in cluster units) * cluster -> block units
    assert native.gridDimX == 4
    assert native.gridDimY == 6
    assert native.gridDimZ == 8
    assert native.blockDimX == 5
    assert native.blockDimY == 6
    assert native.blockDimZ == 7
    assert native.numAttrs == 1
    attr = native.attrs[0]
    assert attr.id == driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION
    assert (attr.value.clusterDim.x, attr.value.clusterDim.y, attr.value.clusterDim.z) == (2, 2, 2)


def test_launch_invalid_values(init_cuda):
    code = 'extern "C" __global__ void my_kernel() {}'
    program = Program(code, SourceCodeType.CXX)
    mod = program.compile(ObjectCodeFormatType.CUBIN)

    stream = Device().create_stream()
    ker = mod.get_kernel("my_kernel")
    config = LaunchConfig(grid=(1, 1, 1), block=(1, 1, 1), shmem_size=0)

    with pytest.raises(TypeError):
        launch(None, ker, config)

    with pytest.raises(TypeError):
        launch(stream, None, config)

    with pytest.raises(TypeError):
        launch(stream, ker, None)

    msg = (
        r"Passing foreign stream objects to this function via the stream "
        r"protocol is deprecated\. Convert the object explicitly using "
        r"Stream\(obj\) instead\."
    )
    with pytest.warns(DeprecationWarning, match=msg):
        launch(StreamWrapper(stream), config, ker)

    launch(stream, config, ker)
    stream.sync()  # TODO(#1539)


# Parametrize: (python_type, cpp_type, init_value)
PARAMS = (
    (bool, "bool", True),
    (float, "double", 2.718),
    (np.bool, "bool", True),
    (np.int8, "signed char", -42),
    (np.int16, "signed short", -1234),
    (np.int32, "signed int", -123456),
    (np.int64, "signed long long", -123456789),
    (np.uint8, "unsigned char", 42),
    (np.uint16, "unsigned short", 1234),
    (np.uint32, "unsigned int", 123456),
    (np.uint64, "unsigned long long", 123456789),
    (np.float32, "float", 3.14),
    (np.float64, "double", 2.718),
    (ctypes.c_bool, "bool", True),
    (ctypes.c_int8, "signed char", -42),
    (ctypes.c_int16, "signed short", -1234),
    (ctypes.c_int32, "signed int", -123456),
    (ctypes.c_int64, "signed long long", -123456789),
    (ctypes.c_uint8, "unsigned char", 42),
    (ctypes.c_uint16, "unsigned short", 1234),
    (ctypes.c_uint32, "unsigned int", 123456),
    (ctypes.c_uint64, "unsigned long long", 123456789),
    (ctypes.c_float, "float", 3.14),
    (ctypes.c_double, "double", 2.718),
)
if helpers.CCCL_INCLUDE_PATHS is not None:
    PARAMS += (
        (np.float16, "half", 0.78),
        (np.complex64, "cuda::std::complex<float>", 1 + 2j),
        (np.complex128, "cuda::std::complex<double>", -3 - 4j),
        (complex, "cuda::std::complex<double>", 5 - 7j),
    )


@pytest.mark.parametrize("python_type, cpp_type, init_value", PARAMS)
@requires_module(np, "2.2.5", reason="need numpy 2.2.5+ (numpy GH #28632)")
def test_launch_scalar_argument(python_type, cpp_type, init_value):
    dev = Device()
    dev.set_current()

    # Prepare pinned host array
    mr = LegacyPinnedMemoryResource()
    b = mr.allocate(np.dtype(python_type).itemsize)
    arr = np.from_dlpack(b).view(python_type)
    arr[:] = 0

    # Prepare scalar argument in Python
    scalar = python_type(init_value)

    # CUDA kernel templated on type T
    code = r"""
    template <typename T>
    __global__ void write_scalar(T* arr, T val) {
        arr[0] = val;
    }
    """

    # Compile and force instantiation for this type
    arch = "".join(f"{i}" for i in dev.compute_capability)
    if helpers.CCCL_INCLUDE_PATHS is not None:
        code = (
            r"""
        #include <cuda_fp16.h>
        #include <cuda/std/complex>
        """
            + code
        )
    pro_opts = ProgramOptions(std="c++17", arch=f"sm_{arch}", include_path=helpers.CCCL_INCLUDE_PATHS)
    prog = Program(code, code_type="c++", options=pro_opts)
    ker_name = f"write_scalar<{cpp_type}>"
    mod = prog.compile("cubin", name_expressions=(ker_name,))
    ker = mod.get_kernel(ker_name)

    # Launch with 1 thread
    stream = dev.default_stream
    config = LaunchConfig(grid=1, block=1)
    launch(stream, config, ker, arr.ctypes.data, scalar)
    stream.sync()

    # Check result
    assert arr[0] == init_value, f"Expected {init_value}, got {arr[0]}"


@skipif_need_cuda_headers  # cg
def test_cooperative_launch():
    dev = Device()
    dev.set_current()
    s = dev.create_stream(options={"nonblocking": True})

    # CUDA kernel templated on type T
    code = r"""
    #include <cooperative_groups.h>

    extern "C" __global__ void test_grid_sync() {
        namespace cg = cooperative_groups;
        auto grid = cg::this_grid();
        grid.sync();
    }
    """

    # Compile and force instantiation for this type
    arch = "".join(f"{i}" for i in dev.compute_capability)
    pro_opts = ProgramOptions(std="c++17", arch=f"sm_{arch}", include_path=helpers.CCCL_INCLUDE_PATHS)
    prog = Program(code, code_type="c++", options=pro_opts)
    ker = prog.compile("cubin").get_kernel("test_grid_sync")

    # # Launch without setting is_cooperative
    # # Commented out as this seems to be a sticky error...
    # config = LaunchConfig(grid=1, block=1)
    # launch(s, config, ker)
    # from cuda.core._utils.cuda_utils import CUDAError
    # with pytest.raises(CUDAError) as e:
    #     s.sync()
    # assert "CUDA_ERROR_LAUNCH_FAILED" in str(e)

    # Crazy grid sizes would not work
    block = 128
    config = LaunchConfig(grid=dev.properties.max_grid_dim_x // block + 1, block=block, is_cooperative=True)
    with pytest.raises(ValueError):
        launch(s, config, ker)

    # This works just fine
    config = LaunchConfig(grid=1, block=1, is_cooperative=True)
    launch(s, config, ker)
    s.sync()


@pytest.mark.skipif(cp is None, reason="cupy not installed")
@pytest.mark.parametrize(
    "memory_resource_class",
    [
        "device_memory_resource",  # kludgy, but can go away after #726 is resolved
        pytest.param(
            LegacyPinnedMemoryResource,
            marks=requires_module(np, "2.2.5", reason="need numpy 2.2.5+ (numpy GH #28632)"),
        ),
    ],
)
def test_launch_with_buffers_allocated_by_memory_resource(init_cuda, memory_resource_class):
    """Test that kernels can access memory allocated by memory resources."""
    dev = Device()
    dev.set_current()
    stream = dev.create_stream()
    # tell CuPy to use our stream as the current stream:
    cp.cuda.ExternalStream(int(stream.handle)).use()

    # Kernel that operates on memory
    code = """
    extern "C"
    __global__ void memory_ops(float* data, size_t N) {
        const unsigned int tid = threadIdx.x + blockIdx.x * blockDim.x;
        if (tid < N) {
            // Access memory (device or pinned)
            data[tid] = data[tid] * 3.0f;
        }
    }
    """

    # Compile kernel
    arch = "".join(f"{i}" for i in dev.compute_capability)
    program_options = ProgramOptions(std="c++17", arch=f"sm_{arch}")
    prog = Program(code, code_type="c++", options=program_options)
    mod = prog.compile("cubin")
    kernel = mod.get_kernel("memory_ops")

    # Create memory resource
    if memory_resource_class == "device_memory_resource":
        if dev.properties.memory_pools_supported:
            mr = DeviceMemoryResource(dev.device_id)
        else:
            mr = _SynchronousMemoryResource(dev.device_id)
        name = memory_resource_class
    else:
        mr = memory_resource_class()
        name = str(mr)

    # Allocate memory
    size = 1024
    dtype = np.float32
    element_size = dtype().itemsize
    total_size = size * element_size

    buffer = mr.allocate(total_size, stream=stream)

    # Create array view based on memory type
    if mr.is_host_accessible:
        # For pinned memory, use numpy
        array = np.from_dlpack(buffer).view(dtype=dtype)
    else:
        array = cp.from_dlpack(buffer).view(dtype=dtype)

    # Initialize data with random values
    if mr.is_host_accessible:
        rng = np.random.default_rng()
        array[:] = rng.random(size, dtype=dtype)
    else:
        rng = cp.random.default_rng()
        array[:] = rng.random(size, dtype=dtype)

    # Store original values for verification
    original = array.copy()

    # Sync before kernel launch
    stream.sync()

    # Launch kernel
    block = 256
    grid = (size + block - 1) // block
    config = LaunchConfig(grid=grid, block=block)

    launch(stream, config, kernel, buffer, np.uint64(size))
    stream.sync()

    # Verify kernel operations
    assert cp.allclose(array, original * 3.0), f"{name} operation failed"

    # Clean up
    buffer.close(stream)
    stream.close()

    cp.cuda.Stream.null.use()  # reset CuPy's current stream to the null stream

    # Verify buffer is properly closed
    assert buffer.handle == 0, f"{name} buffer should be closed"


def test_kernel_arg_unsupported_type():
    from cuda.core._kernel_arg_handler import ParamHolder

    with pytest.raises(TypeError, match="unsupported type"):
        ParamHolder(["not_a_valid_kernel_arg"])


def test_kernel_arg_ctypes_subclass_isinstance_fallback():
    """Subclassed ctypes types hit the isinstance fallback in prepare_ctypes_arg."""
    from cuda.core._kernel_arg_handler import ParamHolder

    class MyInt32(ctypes.c_int32):
        pass

    class MyFloat(ctypes.c_float):
        pass

    class MyBool(ctypes.c_bool):
        pass

    # These should NOT raise; they should be handled via isinstance fallback
    holder = ParamHolder([MyInt32(42), MyFloat(3.14), MyBool(True)])
    assert holder.ptr != 0


@pytest.mark.agent_authored(model="claude-opus-4.8")
@requires_module(np, "2.2.5", reason="need numpy 2.2.5+ (numpy GH #28632)")
@pytest.mark.parametrize(
    ("base_type", "np_dtype", "cpp_type", "raw_value"),
    [
        # ctypes scalar subclasses — one per prepare_ctypes_arg isinstance-fallback
        # branch. Values are chosen to expose a wrong width/sign: unsigned values
        # exceed the same-width signed max, and c_uint64 exceeds uint32 max so a
        # uint64 branch misrouted to prepare_arg[uint32_t] truncates 0x1_0000_0001
        # to 1 and fails the readback.
        (ctypes.c_bool, np.bool_, "bool", True),
        (ctypes.c_int8, np.int8, "signed char", -42),
        (ctypes.c_int16, np.int16, "signed short", -1234),
        (ctypes.c_int32, np.int32, "signed int", -123456),
        (ctypes.c_int64, np.int64, "signed long long", -123456789),
        (ctypes.c_uint8, np.uint8, "unsigned char", 200),
        (ctypes.c_uint16, np.uint16, "unsigned short", 60000),
        (ctypes.c_uint32, np.uint32, "unsigned int", 4000000000),
        (ctypes.c_uint64, np.uint64, "unsigned long long", 0x1_0000_0001),
        (ctypes.c_float, np.float32, "float", 3.14),
        (ctypes.c_double, np.float64, "double", 2.718281828),
        # numpy scalar subclass — prepare_numpy_arg fallback
        (np.float32, np.float32, "float", 3.14),
    ],
    ids=[
        "ctypes_bool",
        "ctypes_int8",
        "ctypes_int16",
        "ctypes_int32",
        "ctypes_int64",
        "ctypes_uint8",
        "ctypes_uint16",
        "ctypes_uint32",
        "ctypes_uint64",
        "ctypes_float",
        "ctypes_double",
        "numpy_float32",
    ],
)
def test_launch_scalar_argument_subclass_fallback(base_type, np_dtype, cpp_type, raw_value):
    """Subclassed scalar arguments survive fallback handling and reach the kernel
    with the correct width/sign. The readback value (not just ptr != 0) guards each
    fallback branch against marshalling the wrong C type, e.g. uint64 -> uint32_t."""

    class Subclassed(base_type):
        pass

    scalar = Subclassed(raw_value)
    expected = np_dtype(raw_value)

    dev = Device()
    dev.set_current()

    mr = LegacyPinnedMemoryResource()
    b = mr.allocate(np.dtype(np_dtype).itemsize)
    arr = np.from_dlpack(b).view(np_dtype)
    arr[:] = 0

    code = r"""
    template <typename T>
    __global__ void write_scalar(T* arr, T val) {
        arr[0] = val;
    }
    """

    arch = "".join(f"{i}" for i in dev.compute_capability)
    pro_opts = ProgramOptions(std="c++17", arch=f"sm_{arch}")
    prog = Program(code, code_type="c++", options=pro_opts)
    ker_name = f"write_scalar<{cpp_type}>"
    mod = prog.compile("cubin", name_expressions=(ker_name,))
    ker = mod.get_kernel(ker_name)

    stream = dev.default_stream
    config = LaunchConfig(grid=1, block=1)
    launch(stream, config, ker, arr.ctypes.data, scalar)
    stream.sync()

    assert arr[0] == expected


def test_kernel_arg_numpy_subclass_isinstance_fallback():
    """Subclassed numpy scalars hit the isinstance fallback in prepare_numpy_arg."""
    from cuda.core._kernel_arg_handler import ParamHolder

    class MyInt32(np.int32):
        pass

    class MyFloat32(np.float32):
        pass

    holder = ParamHolder([MyInt32(7), MyFloat32(2.5)])
    assert holder.ptr != 0


def test_kernel_arg_python_isinstance_fallbacks():
    """Subclassed Python builtins hit the isinstance fallback in ParamHolder."""
    from cuda.core._kernel_arg_handler import ParamHolder

    class MyBool(int):
        """type(x) is not int, so fast path skips; isinstance(x, int) catches it."""

    class MyFloat(float):
        pass

    class MyComplex(complex):
        pass

    holder = ParamHolder([MyBool(1), MyFloat(1.5), MyComplex(1 + 2j)])
    assert holder.ptr != 0
