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

from __future__ import annotations

from libc.stdint cimport uintptr_t, INT32_MIN
from libc.stdlib cimport strtol, getenv

from cuda.bindings cimport cydriver

from cuda.core._event cimport Event as cyEvent
from cuda.core._utils.cuda_utils cimport (
    check_or_create_options,
    HANDLE_RETURN,
)

import cython
import warnings
from dataclasses import dataclass
from typing import Protocol, TYPE_CHECKING

from cuda.core._context cimport Context
from cuda.core._device_resources cimport DeviceResources
from cuda.core._event import Event, EventOptions

from cuda.core._resource_handles cimport (
    ContextHandle,
    EventHandle,
    StreamHandle,
    create_context_handle_ref,
    create_event_handle_noctx,
    create_stream_handle,
    create_stream_handle_with_owner,
    get_current_context,
    get_last_error,
    get_legacy_stream,
    get_per_thread_stream,
    get_stream_context,
    as_intptr,
    as_cu,
    as_py,
)

if TYPE_CHECKING:
    import cuda.bindings.driver  # no-cython-lint
    from cuda.core._device import Device
    from cuda.core.graph import GraphBuilder

__all__ = ['LEGACY_DEFAULT_STREAM', 'PER_THREAD_DEFAULT_STREAM', 'Stream', 'StreamOptions']


@dataclass
cdef class StreamOptions:
    """Customizable :obj:`~_stream.Stream` options.

    Attributes
    ----------
    nonblocking : bool, optional
        Stream does not synchronize with the NULL stream. (Default to True)
    priority : int, optional
        Stream priority where lower number represents a
        higher priority. (Default to lowest priority)

    """

    nonblocking : cython.bint = True
    priority: int | None = None


class IsStreamType(Protocol):
    def __cuda_stream__(self) -> tuple[int, int]:
        """
        For any Python object that is meant to be interpreted as a CUDA stream, the intent
        can be communicated by implementing this protocol that returns a 2-tuple: The protocol
        version number (currently ``0``) and the address of ``cudaStream_t``. Both values
        should be Python `int`.
        """
        ...


cdef class Stream:
    """Represent a queue of GPU operations that are executed in a specific order.

    Applications use streams to control the order of execution for
    GPU work. Work within a single stream are executed sequentially.
    Whereas work across multiple streams can be further controlled
    using stream priorities and :obj:`~_event.Event` managements.

    Advanced users can utilize default streams for enforce complex
    implicit synchronization behaviors.

    Directly creating a :obj:`~_stream.Stream` is not supported due to ambiguity.
    New streams should instead be created through a :obj:`~_device.Device`
    object, or created directly through using an existing handle
    using Stream.from_handle().
    """
    def __init__(self, *args, **kwargs) -> None:
        raise RuntimeError(
            "Stream objects cannot be instantiated directly. "
            "Please use Device APIs (create_stream) or other Stream APIs (from_handle)."
        )

    @staticmethod
    cdef Stream _from_handle(type cls, StreamHandle h_stream):
        """Create a Stream from an existing StreamHandle (cdef-only factory)."""
        cdef Stream s = cls.__new__(cls)
        s._h_stream = h_stream
        s._h_context = get_stream_context(h_stream)
        s._device_id = -1  # lazy init'd (invalid sentinel)
        s._nonblocking = -1  # lazy init'd
        s._priority = INT32_MIN  # lazy init'd
        return s

    @classmethod
    def _legacy_default(cls) -> Stream:
        """Return the legacy default stream (supports subclassing)."""
        return Stream._from_handle(cls, get_legacy_stream())

    @classmethod
    def _per_thread_default(cls) -> Stream:
        """Return the per-thread default stream (supports subclassing)."""
        return Stream._from_handle(cls, get_per_thread_stream())

    @classmethod
    def _init(cls, obj: IsStreamType | None = None, options: object = None,
              device_id: int | None = None, ctx: Context | None = None) -> Stream:
        cdef StreamHandle h_stream
        cdef cydriver.CUstream borrowed
        cdef ContextHandle h_context

        # Extract context handle if provided
        if ctx is not None:
            h_context = (<Context>ctx)._h_context

        if obj is not None and options is not None:
            raise ValueError("obj and options cannot be both specified")
        if obj is not None:
            # Borrowed stream from foreign object
            # C++ handle prevents owner from being GC'd until handle is released
            # Owner is responsible for keeping the stream's context alive
            borrowed = _handle_from_stream_protocol(obj)
            h_stream = create_stream_handle_with_owner(borrowed, obj)
            return Stream._from_handle(cls, h_stream)

        cdef StreamOptions opts = check_or_create_options(StreamOptions, options, "Stream options")
        nonblocking = opts.nonblocking
        priority = opts.priority

        cdef unsigned int flags = (cydriver.CUstream_flags.CU_STREAM_NON_BLOCKING if nonblocking
                                   else cydriver.CUstream_flags.CU_STREAM_DEFAULT)
        # TODO: we might want to consider memoizing high/low per CUDA context and avoid this call
        cdef int high, low
        cdef cydriver.CUresult res_code
        with nogil:
            res_code = cydriver.cuCtxGetStreamPriorityRange(&high, &low)
        if res_code != cydriver.CUresult.CUDA_SUCCESS:
            if res_code == cydriver.CUresult.CUDA_ERROR_INVALID_CONTEXT:
                raise RuntimeError(
                    "No current CUDA context. Call dev.set_current() before creating streams."
                )
            HANDLE_RETURN(res_code)
        cdef int prio
        if priority is not None:
            prio = priority
            if not (low <= prio <= high):
                raise ValueError(f"{priority=} is out of range {[low, high]}")
        else:
            prio = high

        # C++ creates the stream and returns owning handle with context dependency.
        # For green contexts, the C++ layer auto-dispatches to cuGreenCtxStreamCreate.
        h_stream = create_stream_handle(h_context, flags, prio)
        if not h_stream:
            res_code = get_last_error()
            if not nonblocking and res_code == cydriver.CUresult.CUDA_ERROR_INVALID_VALUE:
                # cuGreenCtxStreamCreate rejects CU_STREAM_DEFAULT;
                # no need to check is_green since primary streams don't fail this way
                raise ValueError(
                    "Green context streams must be non-blocking. "
                    "Use StreamOptions(nonblocking=True) or omit the option (True is the default)."
                )
            elif res_code == cydriver.CUresult.CUDA_ERROR_NOT_SUPPORTED:
                raise RuntimeError(
                    "cuGreenCtxStreamCreate is not available. "
                    "Green context stream creation requires CUDA 12.5 or newer."
                )
            else:
                HANDLE_RETURN(res_code)
        cdef Stream self = Stream._from_handle(cls, h_stream)
        self._nonblocking = int(nonblocking)
        self._priority = prio
        if device_id is not None:
            self._device_id = device_id
        return self

    cpdef close(self):
        """Destroy the stream.

        Releases the stream handle. For owned streams, this destroys the
        underlying CUDA stream. For borrowed streams, this releases the
        reference and allows the Python owner to be GC'd.
        """
        self._h_stream.reset()

    def __cuda_stream__(self) -> tuple[int, int]:
        """Return an instance of a __cuda_stream__ protocol."""
        return (0, as_intptr(self._h_stream))

    def __hash__(self) -> int:
        return hash(as_intptr(self._h_stream))

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Stream):
            return NotImplemented
        return as_intptr(self._h_stream) == as_intptr((<Stream>other)._h_stream)

    def __repr__(self) -> str:
        cdef ContextHandle h_context
        Stream_get_ctx(self, &h_context)
        return f"<Stream handle={as_intptr(self._h_stream):#x} context={as_intptr(h_context):#x}>"

    @property
    def handle(self) -> cuda.bindings.driver.CUstream:
        """Return the underlying ``CUstream`` object.

        .. caution::

            This handle is a Python object. To get the memory address of the underlying C
            handle, call ``int(Stream.handle)``.
        """
        return as_py(self._h_stream)

    @property
    def is_nonblocking(self) -> bool:
        """Return True if this is a nonblocking stream, otherwise False."""
        cdef unsigned int flags
        if self._nonblocking == -1:
            with nogil:
                HANDLE_RETURN(cydriver.cuStreamGetFlags(as_cu(self._h_stream), &flags))
            self._nonblocking = flags & cydriver.CUstream_flags.CU_STREAM_NON_BLOCKING
        return bool(self._nonblocking)

    @property
    def priority(self) -> int:
        """Return the stream priority."""
        cdef int prio
        if self._priority == INT32_MIN:
            with nogil:
                HANDLE_RETURN(cydriver.cuStreamGetPriority(as_cu(self._h_stream), &prio))
            self._priority = prio
        return self._priority

    def sync(self) -> None:
        """Synchronize the stream."""
        with nogil:
            HANDLE_RETURN(cydriver.cuStreamSynchronize(as_cu(self._h_stream)))

    def record(self, event: Event | None = None, options: EventOptions | None = None) -> Event:
        """Record an event onto the stream.

        Creates an :obj:`~_event.Event` object (or reuses the given one) by
        recording on the stream.

        Parameters
        ----------
        event : :obj:`~_event.Event`, optional
            Optional event object to be reused for recording.
        options : :obj:`EventOptions`, optional
            Customizable dataclass for event creation options.

        Returns
        -------
        :obj:`~_event.Event`
            Newly created event object.

        Note
        ----
        A default stream (:obj:`LEGACY_DEFAULT_STREAM` or
        :obj:`PER_THREAD_DEFAULT_STREAM`) carries no context of its own; it
        refers to whichever context is current, so a newly created event is
        associated with the current context at call time.

        """
        # Create an Event object (or reusing the given one) by recording
        # on the stream. Event flags such as disabling timing, nonblocking,
        # and CU_EVENT_RECORD_EXTERNAL, can be set in EventOptions.
        cdef ContextHandle h_context
        cdef int device_id
        if event is None:
            Stream_get_ctx_device(self, &h_context, &device_id)
            event = cyEvent._init(cyEvent, device_id, h_context, options, False)
        elif event.is_ipc_enabled:
            raise TypeError(
                "IPC-enabled events should not be re-recorded, instead create a "
                "new event by supplying options."
            )

        cdef cydriver.CUevent e = as_cu((<cyEvent?>(event))._h_event)
        with nogil:
            HANDLE_RETURN(cydriver.cuEventRecord(e, as_cu(self._h_stream)))
        return event

    def wait(self, event_or_stream: Event | Stream) -> None:
        """Wait for a CUDA event or a CUDA stream.

        Waiting for an event or a stream establishes a stream order.

        If a :obj:`~_stream.Stream` is provided, then wait until the stream's
        work is completed. This is done by recording a new :obj:`~_event.Event`
        on the stream and then waiting on it.

        Parameters
        ----------
        event_or_stream : :obj:`~_event.Event` | :obj:`~_stream.Stream`
            The event or stream to wait for. Objects supporting the
            ``__cuda_stream__`` protocol are also accepted and treated as
            streams.

        """
        cdef Stream stream
        cdef EventHandle h_event

        # Handle Event directly
        if isinstance(event_or_stream, Event):
            with nogil:
                # TODO: support flags other than 0?
                HANDLE_RETURN(cydriver.cuStreamWaitEvent(
                    as_cu(self._h_stream), as_cu((<cyEvent>event_or_stream)._h_event), 0))
            return

        # Convert to Stream if needed
        if isinstance(event_or_stream, Stream):
            stream = <Stream>event_or_stream
        else:
            try:
                stream = Stream._init(obj=event_or_stream)
            except Exception as e:
                raise ValueError(
                    "only an Event, Stream, or object supporting __cuda_stream__ can be waited,"
                    f" got {type(event_or_stream)}"
                ) from e

        # Wait on stream via temporary event
        with nogil:
            h_event = create_event_handle_noctx(cydriver.CUevent_flags.CU_EVENT_DISABLE_TIMING)
            HANDLE_RETURN(cydriver.cuEventRecord(as_cu(h_event), as_cu(stream._h_stream)))
            # TODO: support flags other than 0?
            HANDLE_RETURN(cydriver.cuStreamWaitEvent(as_cu(self._h_stream), as_cu(h_event), 0))

    @property
    def device(self) -> Device:
        """Return the :obj:`~_device.Device` singleton associated with this stream.

        Note
        ----
        A default stream (:obj:`LEGACY_DEFAULT_STREAM` or
        :obj:`PER_THREAD_DEFAULT_STREAM`) carries no context of its own; it
        refers to whichever context is current, so this returns the device for
        the current context at call time.

        For a created stream, the current context on the device may differ from
        this stream's context. That case occurs when a different CUDA context is
        set current after the stream is created.

        """
        from cuda.core._device import Device  # avoid circular import
        cdef ContextHandle h_context
        cdef int device_id
        Stream_get_ctx_device(self, &h_context, &device_id)
        return Device(device_id)

    @property
    def context(self) -> Context:
        """Return the :obj:`~_context.Context` associated with this stream.

        Note
        ----
        A default stream (:obj:`LEGACY_DEFAULT_STREAM` or
        :obj:`PER_THREAD_DEFAULT_STREAM`) carries no context of its own; it
        refers to whichever context is current, so this returns the current
        context at call time.

        """
        cdef ContextHandle h_context
        cdef int device_id
        Stream_get_ctx_device(self, &h_context, &device_id)
        return Context._from_handle(Context, h_context, device_id)

    @property
    def resources(self) -> DeviceResources:
        """Query the hardware resources provisioned for this stream's context.

        For streams created from a green context, returns the resources
        that context was provisioned with. For streams on the primary
        context, returns the full device resources.

        Note
        ----
        A default stream (:obj:`LEGACY_DEFAULT_STREAM` or
        :obj:`PER_THREAD_DEFAULT_STREAM`) carries no context of its own; it
        refers to whichever context is current, so this queries the current
        context at call time.

        """
        cdef ContextHandle h_context
        cdef int device_id
        Stream_get_ctx_device(self, &h_context, &device_id)
        return DeviceResources._init_from_ctx(h_context, device_id)

    @staticmethod
    def from_handle(handle) -> Stream:
        """Create a new :obj:`~_stream.Stream` object from a foreign stream handle.

        Uses a cudaStream_t pointer address represented as a Python int
        to create a new :obj:`~_stream.Stream` object.

        Note
        ----
        Stream lifetime is not managed, foreign object must remain
        alive while this stream is active.

        Parameters
        ----------
        handle : int
            Stream handle representing the address of a foreign
            stream object.

        Returns
        -------
        :obj:`~_stream.Stream`
            Newly created stream object.

        """

        class _stream_holder:
            def __cuda_stream__(self) -> tuple[int, int]:
                return (0, handle)

        return Stream._init(obj=_stream_holder())

    def create_graph_builder(self) -> GraphBuilder:
        """Create a new :obj:`~graph.GraphBuilder` object.

        The new graph builder will be associated with this stream.

        Returns
        -------
        :obj:`~graph.GraphBuilder`
            Newly created graph builder object.

        """
        from cuda.core.graph._graph_builder import GraphBuilder

        return GraphBuilder._init(self)


LEGACY_DEFAULT_STREAM: Stream = Stream._legacy_default()
PER_THREAD_DEFAULT_STREAM: Stream = Stream._per_thread_default()


cpdef Stream default_stream():
    """Return the default CUDA :obj:`~_stream.Stream`.

    The type of default stream returned depends on if the environment
    variable CUDA_PYTHON_CUDA_PER_THREAD_DEFAULT_STREAM is set.

    If set, returns a per-thread default stream. Otherwise returns
    the legacy stream.

    """
    # TODO: flip the default
    cdef const char* use_ptds_raw = getenv("CUDA_PYTHON_CUDA_PER_THREAD_DEFAULT_STREAM")

    cdef int use_ptds = 0
    if use_ptds_raw != NULL:
        use_ptds = strtol(use_ptds_raw, NULL, 10)

    # value is non-zero, including for weird stuff like 123foo
    if use_ptds:
        return PER_THREAD_DEFAULT_STREAM
    else:
        return LEGACY_DEFAULT_STREAM


cdef inline bint Stream_is_default_token(Stream self) noexcept nogil:
    """Return True for CU_STREAM_LEGACY and CU_STREAM_PER_THREAD.

    These tokens carry no context of their own; they refer to whatever context
    is current, so nothing resolved from one may be cached on the object.
    """
    cdef uintptr_t h = <uintptr_t>as_cu(self._h_stream)
    return h == <uintptr_t>cydriver.CU_STREAM_LEGACY or h == <uintptr_t>cydriver.CU_STREAM_PER_THREAD


cdef inline int Stream_get_ctx(Stream self, ContextHandle* h_context) except?-1 nogil:
    """Resolve the stream's context handle into ``h_context``."""
    cdef cydriver.CUcontext ctx
    cdef bint is_default = Stream_is_default_token(self)

    # Default-stream tokens must never reuse a sticky object field, even if
    # something else populated ``_h_context`` (defense in depth for #2485).
    if self._h_context and not is_default:
        h_context[0] = self._h_context
        return 0

    h_context[0] = get_stream_context(self._h_stream)
    if not h_context[0]:
        HANDLE_RETURN(cydriver.cuStreamGetCtx(as_cu(self._h_stream), &ctx))
        if ctx != NULL:
            with gil:
                h_context[0] = create_context_handle_ref(ctx)

    if h_context[0] and not is_default:
        self._h_context = h_context[0]
    return 0


cdef inline int Stream_get_ctx_device(Stream self, ContextHandle* h_context, int* device_id) except?-1:
    """Resolve the stream's context handle and device ID."""
    cdef cydriver.CUcontext ctx
    cdef cydriver.CUdevice target_dev
    cdef ContextHandle current_context
    cdef bint switch_context
    cdef bint is_default = Stream_is_default_token(self)

    with nogil:
        Stream_get_ctx(self, h_context)
        if self._device_id >= 0 and not is_default:
            device_id[0] = self._device_id
        else:
            # Get device ID from context, switching context temporarily if needed
            current_context = get_current_context()
            switch_context = (as_cu(current_context) != as_cu(h_context[0]))
            if switch_context:
                HANDLE_RETURN(cydriver.cuCtxPushCurrent(as_cu(h_context[0])))
            HANDLE_RETURN(cydriver.cuCtxGetDevice(&target_dev))
            if switch_context:
                HANDLE_RETURN(cydriver.cuCtxPopCurrent(&ctx))
            device_id[0] = <int>target_dev
            if not is_default:
                self._device_id = device_id[0]
    return 0


cdef cydriver.CUstream _handle_from_stream_protocol(obj) except*:
    if isinstance(obj, Stream):
        return <cydriver.CUstream><uintptr_t>(obj.handle)

    try:
        cuda_stream_attr = obj.__cuda_stream__
    except AttributeError:
        raise TypeError(f"{type(obj)} object does not have a '__cuda_stream__' attribute") from None

    if callable(cuda_stream_attr):
        info = cuda_stream_attr()
    else:
        info = cuda_stream_attr
        warnings.simplefilter("once", DeprecationWarning)
        warnings.warn(
            "Implementing __cuda_stream__ as an attribute is deprecated; it must be implemented as a method",
            stacklevel=3,
            category=DeprecationWarning,
        )

    try:
        len_info = len(info)
    except TypeError as e:
        raise RuntimeError(f"obj.__cuda_stream__ must return a sequence with 2 elements, got {type(info)}") from e
    if len_info != 2:
        raise RuntimeError(f"obj.__cuda_stream__ must return a sequence with 2 elements, got {len_info} elements")
    if info[0] != 0:
        raise RuntimeError(
            f"The first element of the sequence returned by obj.__cuda_stream__ must be 0, got {repr(info[0])}"
        )
    return <cydriver.CUstream><uintptr_t>(info[1])

# Helper for API functions that accept either Stream or GraphBuilder. Performs
# needed checks and returns the relevant stream. Rejects None so that callers
# cannot rely on an implicit fallback to the default stream; if the default
# stream is wanted, pass `device.default_stream` explicitly.
cpdef Stream Stream_accept(arg, bint allow_stream_protocol=False):
    from cuda.core.graph._graph_builder import GraphBuilder

    if arg is None:
        raise TypeError(
            "stream is required and must not be None; "
            "pass device.default_stream explicitly to use the default stream."
        )
    if isinstance(arg, Stream):
        return <Stream>(arg)
    elif isinstance(arg, GraphBuilder):
        return <Stream>(arg.stream)
    elif allow_stream_protocol and hasattr(arg, "__cuda_stream__"):
        stream = Stream._init(arg)
        warnings.warn(
            "Passing foreign stream objects to this function via the "
            "stream protocol is deprecated. Convert the object explicitly "
            "using Stream(obj) instead.",
            stacklevel=2,
            category=DeprecationWarning,
        )
        return <Stream>(stream)
    raise TypeError(f"Stream or GraphBuilder expected, got {type(arg).__name__}")
