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

from __future__ import annotations

from collections.abc import Sequence
from typing import TYPE_CHECKING

IF CUDA_CORE_BUILD_MAJOR >= 13:
    from libcpp.vector cimport vector

from cuda.bindings cimport cydriver
from cuda.core._memory._buffer cimport Buffer
from cuda.core._resource_handles cimport as_cu
from cuda.core._stream cimport Stream, Stream_accept
from cuda.core._utils.cuda_utils cimport HANDLE_RETURN

from cuda.core._host import Host
from cuda.core._utils.cuda_utils import driver
from cuda.core._memory._managed_location import _coerce_location

if TYPE_CHECKING:
    from cuda.core._graph import GraphBuilder
    from cuda.core._device import Device

cdef frozenset _ALL_LOCATION_TYPES = frozenset(("device", "host", "host_numa", "host_numa_current"))
cdef frozenset _DEVICE_HOST_NUMA = frozenset(("device", "host", "host_numa"))
cdef frozenset _DEVICE_HOST_ONLY = frozenset(("device", "host"))

cdef set _ADVICE_IGNORES_LOCATION = {
    driver.CUmem_advise.CU_MEM_ADVISE_SET_READ_MOSTLY,
    driver.CUmem_advise.CU_MEM_ADVISE_UNSET_READ_MOSTLY,
    driver.CUmem_advise.CU_MEM_ADVISE_UNSET_PREFERRED_LOCATION,
}

cdef dict _ADVICE_ALLOWED_LOCTYPES = {
    driver.CUmem_advise.CU_MEM_ADVISE_SET_READ_MOSTLY:          _DEVICE_HOST_NUMA,
    driver.CUmem_advise.CU_MEM_ADVISE_UNSET_READ_MOSTLY:        _DEVICE_HOST_NUMA,
    driver.CUmem_advise.CU_MEM_ADVISE_SET_PREFERRED_LOCATION:   _ALL_LOCATION_TYPES,
    driver.CUmem_advise.CU_MEM_ADVISE_UNSET_PREFERRED_LOCATION: _DEVICE_HOST_NUMA,
    driver.CUmem_advise.CU_MEM_ADVISE_SET_ACCESSED_BY:          _DEVICE_HOST_ONLY,
    driver.CUmem_advise.CU_MEM_ADVISE_UNSET_ACCESSED_BY:        _DEVICE_HOST_ONLY,
}


cdef void _require_managed_buffer(Buffer self, str what):
    # Buffer.is_managed handles both pointer-attribute and memory-resource
    # paths (e.g. pool-allocated managed memory whose pointer attribute
    # does not advertise CU_POINTER_ATTRIBUTE_IS_MANAGED).
    if not self.is_managed:
        raise ValueError(f"{what} requires a managed-memory allocation")


cdef tuple _coerce_batch_buffers(object buffers, str what):
    """Coerce ``buffers`` to a tuple[Buffer, ...]; rejects a single Buffer.

    For single-buffer operations, use the corresponding ManagedBuffer
    instance method instead.
    """
    cdef Buffer buf
    cdef list out
    if isinstance(buffers, Buffer):
        raise TypeError(
            f"{what}: pass a sequence of Buffers; for a single buffer use "
            f"the ManagedBuffer instance method"
        )
    if isinstance(buffers, Sequence):
        if not buffers:
            raise ValueError(f"{what}: empty buffers sequence")
        out = []
        for t in buffers:
            buf = <Buffer?>t
            out.append(buf)
        return tuple(out)
    raise TypeError(
        f"{what}: buffers must be a sequence of Buffer, "
        f"got {type(buffers).__name__}"
    )


cdef tuple _broadcast_locations(object location, Py_ssize_t n, bint allow_none, str what):
    if isinstance(location, Sequence):
        if len(location) != n:
            raise ValueError(
                f"{what}: location length {len(location)} does not match "
                f"targets length {n}"
            )
        return tuple(_coerce_location(loc, allow_none=allow_none) for loc in location)
    cdef object coerced = _coerce_location(location, allow_none=allow_none)
    return tuple([coerced] * n)


IF CUDA_CORE_BUILD_MAJOR >= 13:
    # Convert a _LocSpec dataclass to a cydriver.CUmemLocation struct.
    cdef inline cydriver.CUmemLocation _to_cumemlocation(object loc):
        cdef str kind = loc.kind
        if kind == "device":
            return cydriver.CUmemLocation(
                type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE,
                id=<int>loc.id)
        elif kind == "host":
            return cydriver.CUmemLocation(
                type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST,
                id=0)
        elif kind == "host_numa":
            return cydriver.CUmemLocation(
                type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA,
                id=<int>loc.id)
        else:  # host_numa_current
            return cydriver.CUmemLocation(
                type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT,
                id=0)
ELSE:
    # CUDA 12 cuMemPrefetchAsync takes a device ordinal (-1 = host).
    cdef inline int _to_legacy_device(object loc) except? -2:
        cdef str kind = loc.kind
        if kind == "device":
            return <int>loc.id
        if kind == "host":
            return -1
        raise RuntimeError(
            "Host(numa_id=...) / Host.numa_current() require both cuda-bindings 13.0+ "
            "and a CUDA 13+ runtime driver; use Host() instead"
        )


def discard_batch(stream: Stream | GraphBuilder, buffers: Sequence[Buffer]) -> None:
    """Discard a batch of managed-memory ranges.

    Requires CUDA 13+. For a single buffer, use
    :meth:`ManagedBuffer.discard` instead.

    Parameters
    ----------
    stream : :class:`~_stream.Stream` | :class:`~graph.GraphBuilder`
        Stream for the asynchronous discard. First positional, required
        (mirrors :func:`launch`).
    buffers : Sequence[:class:`Buffer`]
        Two or more managed allocations to discard. Resident pages are
        released without prefetching new contents; subsequent access is
        satisfied by lazy migration.

    Raises
    ------
    NotImplementedError
        On a CUDA 12 build of ``cuda.core``.
    """
    cdef tuple bufs = _coerce_batch_buffers(buffers, "discard_batch")
    cdef Stream s = Stream_accept(stream)

    cdef Buffer buf
    for buf in bufs:
        _require_managed_buffer(buf, "discard_batch")

    _do_batch_discard(bufs, s)


def _do_single_discard_py(Buffer buf, stream: Stream | GraphBuilder | None) -> None:
    """Internal: single-buffer discard for ManagedBuffer.discard()."""
    _require_managed_buffer(buf, "discard")
    cdef Stream s = Stream_accept(stream)
    # No single-range cuMemDiscard exists; route through the batched call
    # with count=1.
    cdef tuple bufs = (buf,)
    _do_batch_discard(bufs, s)


cdef void _do_batch_discard(tuple bufs, Stream s):
    IF CUDA_CORE_BUILD_MAJOR >= 13:
        cdef Py_ssize_t n = len(bufs)
        cdef cydriver.CUstream hstream = as_cu(s._h_stream)
        cdef vector[cydriver.CUdeviceptr] ptrs
        cdef vector[size_t] sizes
        ptrs.resize(n)
        sizes.resize(n)
        cdef Buffer buf
        cdef Py_ssize_t i
        for i in range(n):
            buf = <Buffer>bufs[i]
            ptrs[i] = as_cu(buf._h_ptr)
            sizes[i] = buf._size
        with nogil:
            HANDLE_RETURN(cydriver.cuMemDiscardBatchAsync(
                ptrs.data(), sizes.data(), <size_t>n, 0, hstream,
            ))
    ELSE:
        raise NotImplementedError(
            "discard requires a CUDA 13 build of cuda.core"
        )


def _advise_one(Buffer buf, advice: driver.CUmem_advise, location: Device | Host | None) -> None:
    """Internal: apply managed-memory advice to a single buffer.

    Used by :class:`ManagedBuffer` property setters. Not part of the
    public API.
    """
    _require_managed_buffer(buf, "advise")
    if not isinstance(advice, driver.CUmem_advise):
        raise TypeError(
            f"advice must be a cuda.bindings.driver.CUmem_advise value, "
            f"got {type(advice).__name__}"
        )
    cdef frozenset allowed_kinds = _ADVICE_ALLOWED_LOCTYPES.get(advice)
    if allowed_kinds is None:
        raise ValueError(f"Unsupported advice value: {advice!r}")
    cdef bint allow_none = advice in _ADVICE_IGNORES_LOCATION
    cdef object loc = _coerce_location(location, allow_none=allow_none)
    if loc is not None and loc.kind not in allowed_kinds:
        raise ValueError(
            f"advise {advice.name} does not support location_type='{loc.kind}'"
        )
    _do_single_advise(buf, advice, loc, allow_none)


cdef void _do_single_advise(Buffer buf, object advice_value, object loc, bint allow_none):
    cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr)
    cdef size_t nbytes = buf._size
    cdef cydriver.CUmem_advise advice_enum = <cydriver.CUmem_advise>(<int>int(advice_value))
    IF CUDA_CORE_BUILD_MAJOR >= 13:
        cdef cydriver.CUmemLocation cu_loc
        if loc is None:
            # Driver ignores location for read_mostly / unset_preferred_location
            # advice values but still validates the CUmemLocation; pass a
            # host placeholder.
            cu_loc = cydriver.CUmemLocation(
                type=cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST,
                id=0)
        else:
            cu_loc = _to_cumemlocation(loc)
        with nogil:
            HANDLE_RETURN(cydriver.cuMemAdvise(cu_ptr, nbytes, advice_enum, cu_loc))
    ELSE:
        cdef int dev_int = -1 if loc is None else _to_legacy_device(loc)
        with nogil:
            HANDLE_RETURN(cydriver.cuMemAdvise(cu_ptr, nbytes, advice_enum, dev_int))


def prefetch_batch(
    stream: Stream | GraphBuilder,
    buffers: Sequence[Buffer],
    locations: Device | Host | Sequence[Device | Host],
) -> None:
    """Prefetch a batch of managed-memory ranges to target locations.

    Requires CUDA 13+. For a single buffer, use
    :meth:`ManagedBuffer.prefetch` instead.

    Parameters
    ----------
    stream : :class:`~_stream.Stream` | :class:`~graph.GraphBuilder`
        Stream for the asynchronous prefetch. First positional, required
        (mirrors :func:`launch`).
    buffers : Sequence[:class:`Buffer`]
        Two or more managed allocations to operate on.
    locations : :class:`~cuda.core.Device` | :class:`~cuda.core.Host` | Sequence[...]
        Target location(s). A single location applies to all buffers; a
        sequence must match ``len(buffers)``.

    Notes
    -----
    On a CUDA 12 build, falls back to a Python-level loop calling
    ``cuMemPrefetchAsync`` per buffer (no batched driver entry point on
    CUDA 12). CUDA 13 builds use ``cuMemPrefetchBatchAsync`` directly.
    """
    cdef tuple bufs = _coerce_batch_buffers(buffers, "prefetch_batch")
    cdef Py_ssize_t n = len(bufs)
    cdef tuple locs = _broadcast_locations(locations, n, False, "prefetch_batch")
    cdef Stream s = Stream_accept(stream)

    cdef Buffer buf
    for buf in bufs:
        _require_managed_buffer(buf, "prefetch_batch")

    _do_batch_prefetch(bufs, locs, s)


def _do_single_prefetch_py(Buffer buf, location: Device | Host | None, stream: Stream | GraphBuilder | None) -> None:
    """Internal: single-buffer prefetch for ManagedBuffer.prefetch().

    Uses cuMemPrefetchAsync (works on CUDA 12 and 13).
    """
    _require_managed_buffer(buf, "prefetch")
    cdef object loc = _coerce_location(location, allow_none=False)
    cdef Stream s = Stream_accept(stream)
    _do_single_prefetch(buf, loc, s)


cdef void _do_single_prefetch(Buffer buf, object loc, Stream s):
    cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr)
    cdef size_t nbytes = buf._size
    cdef cydriver.CUstream hstream = as_cu(s._h_stream)
    IF CUDA_CORE_BUILD_MAJOR >= 13:
        cdef cydriver.CUmemLocation cu_loc = _to_cumemlocation(loc)
        with nogil:
            HANDLE_RETURN(cydriver.cuMemPrefetchAsync(cu_ptr, nbytes, cu_loc, 0, hstream))
    ELSE:
        cdef int dev_int = _to_legacy_device(loc)
        with nogil:
            HANDLE_RETURN(cydriver.cuMemPrefetchAsync(cu_ptr, nbytes, dev_int, hstream))


IF CUDA_CORE_BUILD_MAJOR >= 13:
    # Function-pointer type for cuMemPrefetchBatchAsync /
    # cuMemDiscardAndPrefetchBatchAsync; both have identical signatures.
    ctypedef cydriver.CUresult (*_BatchPrefetchFn)(
        cydriver.CUdeviceptr*, size_t*, size_t,
        cydriver.CUmemLocation*, size_t*, size_t,
        unsigned long long, cydriver.CUstream,
    ) except ?cydriver.CUDA_ERROR_NOT_FOUND nogil


    def _read_preferred_location_v2(Buffer buf) -> Device | Host | None:
        """Internal: read preferred_location with full NUMA detail.

        Bypasses cuda.bindings.driver.cuMemRangeGetAttribute (whose
        attribute allowlist doesn't yet include the cu13 _TYPE / _ID
        attributes) by calling cydriver directly.

        Returns Device | Host | None.
        """
        cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr)
        cdef size_t nbytes = buf._size
        cdef int loc_type = 0
        cdef int loc_id = 0
        with nogil:
            HANDLE_RETURN(cydriver.cuMemRangeGetAttribute(
                <void*>&loc_type, sizeof(int),
                cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_TYPE,
                cu_ptr, nbytes,
            ))
            HANDLE_RETURN(cydriver.cuMemRangeGetAttribute(
                <void*>&loc_id, sizeof(int),
                cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_ID,
                cu_ptr, nbytes,
            ))
        if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE:
            from cuda.core._device import Device
            return Device(loc_id)
        if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST:
            return Host()
        if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA:
            return Host(numa_id=loc_id)
        if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT:
            return Host.numa_current()
        return None  # CU_MEM_LOCATION_TYPE_INVALID — no preferred location


    cdef void _do_batch_prefetch_op(tuple bufs, tuple locs, Stream s, _BatchPrefetchFn fn):
        """Shared body for batched prefetch / discard-and-prefetch."""
        cdef Py_ssize_t n = len(bufs)
        cdef cydriver.CUstream hstream = as_cu(s._h_stream)
        cdef vector[cydriver.CUdeviceptr] ptrs
        cdef vector[size_t] sizes
        cdef vector[cydriver.CUmemLocation] loc_arr
        cdef vector[size_t] loc_indices
        ptrs.resize(n)
        sizes.resize(n)
        loc_arr.resize(n)
        loc_indices.resize(n)
        cdef Buffer buf
        cdef Py_ssize_t i
        for i in range(n):
            buf = <Buffer>bufs[i]
            ptrs[i] = as_cu(buf._h_ptr)
            sizes[i] = buf._size
            loc_arr[i] = _to_cumemlocation(locs[i])
            loc_indices[i] = <size_t>i
        with nogil:
            HANDLE_RETURN(fn(
                ptrs.data(), sizes.data(), <size_t>n,
                loc_arr.data(), loc_indices.data(), <size_t>n,
                0, hstream,
            ))
ELSE:
    def _read_preferred_location_v2(Buffer buf) -> Device | Host | None:
        # Symbol exists so _managed_buffer.py can `from ... import
        # _read_preferred_location_v2` unconditionally at module top.
        # `ManagedBuffer.preferred_location` gates on both
        # binding_version() and driver_version() >= (13, 0, 0) before
        # calling, so this path is unreachable on a cu12 build.
        raise NotImplementedError(
            "_read_preferred_location_v2 requires a CUDA 13 build of cuda.core"
        )


cdef void _do_batch_prefetch(tuple bufs, tuple locs, Stream s):
    IF CUDA_CORE_BUILD_MAJOR >= 13:
        _do_batch_prefetch_op(bufs, locs, s, cydriver.cuMemPrefetchBatchAsync)
    ELSE:
        # cu12 has no cuMemPrefetchBatchAsync; loop per-range.
        cdef Buffer buf
        cdef Py_ssize_t i
        cdef Py_ssize_t n = len(bufs)
        for i in range(n):
            buf = <Buffer>bufs[i]
            _do_single_prefetch(buf, locs[i], s)


def discard_prefetch_batch(
    stream: Stream | GraphBuilder,
    buffers: Sequence[Buffer],
    locations: Device | Host | Sequence[Device | Host],
) -> None:
    """Discard a batch of managed-memory ranges and prefetch them to target locations.

    Requires CUDA 13+. For a single buffer, use
    :meth:`ManagedBuffer.discard_prefetch` instead.

    Parameters
    ----------
    stream : :class:`~_stream.Stream` | :class:`~graph.GraphBuilder`
        Stream for the asynchronous operation. First positional, required
        (mirrors :func:`launch`).
    buffers : Sequence[:class:`Buffer`]
        Two or more managed allocations to discard and re-prefetch.
    locations : :class:`~cuda.core.Device` | :class:`~cuda.core.Host` | Sequence[...]
        Target location(s). A single location applies to all buffers;
        a sequence must match ``len(buffers)``.

    Raises
    ------
    NotImplementedError
        On a CUDA 12 build of ``cuda.core``.
    """
    cdef tuple bufs = _coerce_batch_buffers(buffers, "discard_prefetch_batch")
    cdef Py_ssize_t n = len(bufs)
    cdef tuple locs = _broadcast_locations(locations, n, False, "discard_prefetch_batch")
    cdef Stream s = Stream_accept(stream)

    cdef Buffer buf
    for buf in bufs:
        _require_managed_buffer(buf, "discard_prefetch_batch")

    _do_batch_discard_prefetch(bufs, locs, s)


def _do_single_discard_prefetch_py(Buffer buf, location: Device | Host | None, stream: Stream | GraphBuilder | None) -> None:
    """Internal: single-buffer discard+prefetch for
    ManagedBuffer.discard_prefetch()."""
    _require_managed_buffer(buf, "discard_prefetch")
    cdef object loc = _coerce_location(location, allow_none=False)
    cdef Stream s = Stream_accept(stream)
    cdef tuple bufs = (buf,)
    cdef tuple locs = (loc,)
    _do_batch_discard_prefetch(bufs, locs, s)


cdef void _do_batch_discard_prefetch(tuple bufs, tuple locs, Stream s):
    IF CUDA_CORE_BUILD_MAJOR >= 13:
        _do_batch_prefetch_op(bufs, locs, s, cydriver.cuMemDiscardAndPrefetchBatchAsync)
    ELSE:
        raise NotImplementedError(
            "discard_prefetch requires a CUDA 13 build of cuda.core"
        )
