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

from dataclasses import dataclass
from typing import TYPE_CHECKING

from libc.stdint cimport intptr_t

from cuda.bindings cimport cydriver

from cuda.core.graph._graph_definition cimport GraphCondition, GraphDefinition
from cuda.core.graph._graph_node cimport GraphNode
from cuda.core.graph._host_callback cimport _resolve_host_callback
from cuda.core.graph._subclasses cimport (
    ExecutableGraphNode,
    create_executable_node_view,
)
from cuda.core._resource_handles cimport (
    GraphExecHandle,
    GraphHandle,
    OpaqueHandle,
    PreparedAttachment,
    as_cu, as_py,
    create_child_graph_handle, create_graph_exec_handle, create_graph_handle,
    get_last_error,
    graph_clone_attachments,
    graph_commit_attachment,
    graph_exec_update,
    graph_prepare_attachment,
    invalidate_child_graph_state,
    retry_deferred_cleanup,
)
from cuda.core._stream cimport Stream
from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
from cuda.core._utils.version cimport cy_binding_version, cy_driver_version

from cuda.core._utils.cuda_utils import (
    CUDAError,
    driver,
    handle_return,
)

if TYPE_CHECKING:
    from cuda.core.graph._graph_definition import GraphDefinition

__all__ = ['Graph', 'GraphBuilder', 'GraphCompleteOptions', 'GraphDebugPrintOptions']


@dataclass
class GraphDebugPrintOptions:
    """Options for debug_dot_print().

    Attributes
    ----------
    verbose : bool
        Output all debug data as if every debug flag is enabled (Default to False)
    runtime_types : bool
        Use CUDA Runtime structures for output (Default to False)
    kernel_node_params : bool
        Adds kernel parameter values to output (Default to False)
    memcpy_node_params : bool
        Adds memcpy parameter values to output (Default to False)
    memset_node_params : bool
        Adds memset parameter values to output (Default to False)
    host_node_params : bool
        Adds host parameter values to output (Default to False)
    event_node_params : bool
        Adds event parameter values to output (Default to False)
    ext_semas_signal_node_params : bool
        Adds external semaphore signal parameter values to output (Default to False)
    ext_semas_wait_node_params : bool
        Adds external semaphore wait parameter values to output (Default to False)
    kernel_node_attributes : bool
        Adds kernel node attributes to output (Default to False)
    handles : bool
        Adds node handles and every kernel function handle to output (Default to False)
    mem_alloc_node_params : bool
        Adds memory alloc parameter values to output (Default to False)
    mem_free_node_params : bool
        Adds memory free parameter values to output (Default to False)
    batch_mem_op_node_params : bool
        Adds batch mem op parameter values to output (Default to False)
    extra_topo_info : bool
        Adds edge numbering information (Default to False)
    conditional_node_params : bool
        Adds conditional node parameter values to output (Default to False)

    """

    verbose: bool = False
    runtime_types: bool = False
    kernel_node_params: bool = False
    memcpy_node_params: bool = False
    memset_node_params: bool = False
    host_node_params: bool = False
    event_node_params: bool = False
    ext_semas_signal_node_params: bool = False
    ext_semas_wait_node_params: bool = False
    kernel_node_attributes: bool = False
    handles: bool = False
    mem_alloc_node_params: bool = False
    mem_free_node_params: bool = False
    batch_mem_op_node_params: bool = False
    extra_topo_info: bool = False
    conditional_node_params: bool = False

    def _to_flags(self) -> int:
        """Convert options to CUDA driver API flags (internal use)."""
        flags = 0
        if self.verbose:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_VERBOSE
        if self.runtime_types:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_RUNTIME_TYPES
        if self.kernel_node_params:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_KERNEL_NODE_PARAMS
        if self.memcpy_node_params:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_MEMCPY_NODE_PARAMS
        if self.memset_node_params:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_MEMSET_NODE_PARAMS
        if self.host_node_params:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_HOST_NODE_PARAMS
        if self.event_node_params:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_EVENT_NODE_PARAMS
        if self.ext_semas_signal_node_params:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_EXT_SEMAS_SIGNAL_NODE_PARAMS
        if self.ext_semas_wait_node_params:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_EXT_SEMAS_WAIT_NODE_PARAMS
        if self.kernel_node_attributes:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_KERNEL_NODE_ATTRIBUTES
        if self.handles:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_HANDLES
        if self.mem_alloc_node_params:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_MEM_ALLOC_NODE_PARAMS
        if self.mem_free_node_params:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_MEM_FREE_NODE_PARAMS
        if self.batch_mem_op_node_params:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_BATCH_MEM_OP_NODE_PARAMS
        if self.extra_topo_info:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_EXTRA_TOPO_INFO
        if self.conditional_node_params:
            flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_CONDITIONAL_NODE_PARAMS
        return flags


@dataclass
class GraphCompleteOptions:
    """Options for graph instantiation.

    Attributes
    ----------
    auto_free_on_launch : bool, optional
        Automatically free memory allocated in a graph before relaunching. (Default to False)
    upload_stream : Stream, optional
        Stream to use to automatically upload the graph after completion. (Default to None)
    device_launch : bool, optional
        Configure the graph to be launchable from the device. This flag can only
        be used on platforms which support unified addressing. This flag cannot be
        used in conjunction with auto_free_on_launch. (Default to False)
    use_node_priority : bool, optional
        Run the graph using the per-node priority attributes rather than the
        priority of the stream it is launched into. (Default to False)

    """

    auto_free_on_launch: bool = False
    upload_stream: Stream | None = None
    device_launch: bool = False
    use_node_priority: bool = False


def _instantiate_graph(source, options: GraphCompleteOptions | None = None) -> Graph:
    cdef GraphHandle h_graph
    cdef GraphExecHandle h_exec

    if isinstance(source, GraphBuilder):
        h_graph = (<GraphBuilder>source)._h_graph
    elif isinstance(source, GraphDefinition):
        h_graph = (<GraphDefinition>source)._h_graph
    else:
        raise TypeError(
            f"expected GraphBuilder or GraphDefinition, got {type(source).__name__}")

    cdef cydriver.CUDA_GRAPH_INSTANTIATE_PARAMS params = cydriver.CUDA_GRAPH_INSTANTIATE_PARAMS(
        flags=0,
        hUploadStream=<cydriver.CUstream>NULL,
        hErrNode_out=<cydriver.CUgraphNode>NULL,
        result_out=cydriver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_SUCCESS,
    )
    if options:
        flags = 0
        if options.auto_free_on_launch:
            flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
        if options.upload_stream:
            flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD
            params.hUploadStream = as_cu((<Stream>options.upload_stream)._h_stream)
        if options.device_launch:
            flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_DEVICE_LAUNCH
        if options.use_node_priority:
            flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
        params.flags = flags

    # The exec is adopted only when result_out reports success, so the
    # diagnostics below run before the handle is checked.
    h_exec = create_graph_exec_handle(h_graph, &params)
    if params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_ERROR:
        raise RuntimeError(
            "Instantiation failed for an unexpected reason which is described in the return value of the function."
        )
    elif params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_INVALID_STRUCTURE:
        raise RuntimeError("Instantiation failed due to invalid structure, such as cycles.")
    elif params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_NODE_OPERATION_NOT_SUPPORTED:
        raise RuntimeError(
            "Instantiation for device launch failed because the graph contained an unsupported operation."
        )
    elif params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_MULTIPLE_CTXS_NOT_SUPPORTED:
        raise RuntimeError("Instantiation for device launch failed due to the nodes belonging to different contexts.")
    elif (
        cy_binding_version() >= (12, 8, 0)
        and params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_CONDITIONAL_HANDLE_UNUSED
    ):
        raise RuntimeError("One or more conditional handles are not associated with conditional builders.")
    elif params.result_out != driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_SUCCESS:
        raise RuntimeError(f"Graph instantiation failed with unexpected error code: {params.result_out}")

    if as_cu(h_exec) == NULL:
        HANDLE_RETURN(get_last_error())
    return Graph._init(h_exec)


# Distinguishes the three kinds of GraphBuilder, which differ in how they
# begin/end stream capture and whether they own the resulting CUgraph.
# Each kind progresses through _CaptureState as follows:
#
#   PRIMARY:          NOT_STARTED -> CAPTURING -> ENDED
#   FORKED:           CAPTURING (never transitions; joined and closed)
#   CONDITIONAL_BODY: NOT_STARTED -> CAPTURING -> ENDED
#
cdef enum _BuilderKind:
    # PRIMARY: The top-level builder created by Device or Stream. Owns the
    # captured CUgraph via an owning GraphHandle. Progresses through all three
    # capture states; responsible for ending capture if destroyed early.
    PRIMARY = 0
    # FORKED: Created by split(). Captures on a private stream forked from the
    # primary. Starts in CAPTURING state and never transitions; the user joins
    # it back to the primary via join(), which closes the builder.  Must NOT
    # call cuStreamEndCapture (the driver requires all forked streams to be
    # joined first).
    FORKED = 1
    # CONDITIONAL_BODY: Created by if_then/if_else/switch/while_loop. Captures
    # into a non-owned body graph via cuStreamBeginCaptureToGraph. The body
    # graph's lifetime is tied to a parent graph. Progresses through all three
    # capture states like PRIMARY.
    CONDITIONAL_BODY = 2


# Tracks the capture lifecycle of a GraphBuilder.
cdef enum _CaptureState:
    CAPTURE_NOT_STARTED = 0
    CAPTURING = 1
    CAPTURE_ENDED = 2  # Finished, valid handle
    CLOSED = 3         # No valid handle


cdef class GraphBuilder:
    """A graph under construction by stream capture.

    A graph groups a set of CUDA kernels and other CUDA operations together and executes
    them with a specified dependency tree. It speeds up the workflow by combining the
    driver activities associated with CUDA kernel launches and CUDA API calls.

    Directly creating a :obj:`~graph.GraphBuilder` is not supported due
    to ambiguity. New graph builders should instead be created through a
    :obj:`~_device.Device`, or a :obj:`~_stream.stream` object.

    .. note::

        Operations recorded during capture reference your memory but do not
        take ownership of it. As with ordinary stream work, you must keep the
        operands alive for as long as the completed graph may execute -- for
        example, the :obj:`~_memory.Buffer` objects passed to :func:`~launch`
        or :meth:`~_memory.Buffer.copy_to`. Host callbacks added with
        :meth:`callback` are the exception: the callable (and any copied
        ``user_data``) are retained for the graph's lifetime. This differs from
        building a graph explicitly with :class:`~graph.GraphDefinition`, which
        retains the operands it is given.

    """

    def __init__(self):
        raise NotImplementedError(
            "directly creating a GraphBuilder object can be ambiguous. Please either "
            "call Device.create_graph_builder() or stream.create_graph_builder()"
        )

    def __dealloc__(self):
        GB_end_capture_if_needed(self, False)

    @staticmethod
    def _init(Stream stream):
        cdef GraphBuilder self = GraphBuilder.__new__(GraphBuilder)
        # _h_graph set by begin_building
        self._h_stream = stream._h_stream
        self._kind = PRIMARY
        self._state = CAPTURE_NOT_STARTED
        self._stream = stream
        return self

    def close(self):
        """Destroy the graph builder."""
        GB_end_capture_if_needed(self, True)
        self._h_graph.reset()
        self._h_stream.reset()
        retry_deferred_cleanup()
        self._state = CLOSED
        self._stream = None

    @property
    def stream(self) -> Stream:
        """Returns the stream associated with the graph builder."""
        return self._stream

    @property
    def is_join_required(self) -> bool:
        """Returns True if this graph builder must be joined before building is ended."""
        return self._kind == FORKED

    @property
    def graph_definition(self) -> GraphDefinition:
        """The captured graph as an explicit :class:`~graph.GraphDefinition`.

        .. versionadded:: 1.1.0

        The returned :class:`~graph.GraphDefinition` is a view of the same
        graph this builder is producing: nodes added through it appear in
        subsequent :meth:`complete` and :meth:`debug_dot_print` calls, and
        the view stays valid even after the builder is closed.

        This lets you mix the capture and explicit APIs on a single graph,
        for example to inspect what was captured, augment it with extra
        nodes, or build a conditional body entirely with the explicit API.

        Availability:

        - **Primary builders** (created by :meth:`Device.create_graph_builder`
          or :meth:`Stream.create_graph_builder`): only after
          :meth:`end_building`.

        - **Conditional-body builders** (returned by :meth:`if_then`,
          :meth:`if_else`, :meth:`while_loop`, :meth:`switch`): both before
          :meth:`begin_building` and after :meth:`end_building`. The body
          graph already exists when the conditional is created, so you may
          populate it through this view without ever calling
          :meth:`begin_building` on the body builder.

        - **Forked builders** (returned by :meth:`split`): never. Forked
          builders share the primary builder's graph; access it through the
          primary instead.

        Returns
        -------
        GraphDefinition
            A view of the graph being built.

        Raises
        ------
        RuntimeError
            If the builder is closed, forked, currently building, or (for
            primary builders) has not started building yet. A
            :class:`~graph.GraphDefinition` obtained before :meth:`close`
            keeps working; only fresh access through this property is
            rejected once the builder is closed.
        """
        GB_check_open(self)
        if self._kind == FORKED:
            raise RuntimeError(
                "graph_definition is unavailable on forked graph builders; "
                "access it through the primary builder instead."
            )
        elif self._state == CAPTURING:
            raise RuntimeError(
                "graph_definition is unavailable while capture is in "
                "progress; call end_building() first."
            )
        elif self._kind == PRIMARY:
            if self._state == CAPTURE_NOT_STARTED:
                raise RuntimeError(
                    "graph_definition is unavailable before begin_building() on "
                    "a primary builder; no graph has been created yet."
                )
        return GraphDefinition._from_handle(self._h_graph)

    def begin_building(self, mode: str | None = "relaxed") -> GraphBuilder:
        """Begins the building process.

        Build `mode` for controlling interaction with other API calls must be one of the following:

        - `global` : Prohibit potentially unsafe operations across all streams in the process.
        - `thread_local` : Prohibit potentially unsafe operations in streams created by the current thread.
        - `relaxed` : The local thread is not prohibited from potentially unsafe operations.

        Parameters
        ----------
        mode : str, optional
            Build mode to control the interaction with other API calls that are porentially unsafe.
            Default set to use relaxed.

        """
        GB_check_open(self)
        if self._state != CAPTURE_NOT_STARTED:
            if self._state == CAPTURING:
                raise RuntimeError("Graph builder is already building.")
            else:
                raise RuntimeError("Cannot resume building after building has ended.")
        cdef cydriver.CUstreamCaptureMode c_mode
        if mode == "global":
            c_mode = cydriver.CU_STREAM_CAPTURE_MODE_GLOBAL
        elif mode == "thread_local":
            c_mode = cydriver.CU_STREAM_CAPTURE_MODE_THREAD_LOCAL
        elif mode == "relaxed":
            c_mode = cydriver.CU_STREAM_CAPTURE_MODE_RELAXED
        else:
            raise ValueError(f"Unsupported build mode: {mode}")

        cdef cydriver.CUstream c_stream = as_cu(self._h_stream)
        cdef cydriver.CUgraph c_graph
        cdef cydriver.CUstreamCaptureStatus c_status
        if self._kind == CONDITIONAL_BODY:
            c_graph = as_cu(self._h_graph)
            with nogil:
                HANDLE_RETURN(cydriver.cuStreamBeginCaptureToGraph(
                    c_stream, c_graph, NULL, NULL, 0, c_mode))
            self._state = CAPTURING
        else:
            with nogil:
                HANDLE_RETURN(cydriver.cuStreamBeginCapture(c_stream, c_mode))
            # Capture is active now; set CAPTURING before the calls below so a
            # failure in _get_capture_info/create_graph_handle still lets
            # cleanup end the capture rather than leaving the stream poisoned.
            self._state = CAPTURING
            with nogil:
                # The driver rejects a NULL captureStatus_out, so pass a
                # stack-local even though we only want the graph handle.
                _get_capture_info(c_stream, &c_status, &c_graph)
            self._h_graph = create_graph_handle(c_graph)
        return self

    @property
    def is_building(self) -> bool:
        """Returns True if the graph builder is currently building."""
        GB_check_open(self)
        cdef cydriver.CUstream c_stream = as_cu(self._h_stream)
        cdef cydriver.CUstreamCaptureStatus status
        with nogil:
            _get_capture_info(c_stream, &status, NULL)
        if status == cydriver.CU_STREAM_CAPTURE_STATUS_NONE:
            return False
        elif status == cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE:
            return True
        elif status == cydriver.CU_STREAM_CAPTURE_STATUS_INVALIDATED:
            raise RuntimeError(
                "Build process encountered an error and has been invalidated. Build process must now be ended."
            )
        else:
            raise NotImplementedError(f"Unsupported capture status type received: {status}")

    def end_building(self) -> GraphBuilder:
        """Ends the building process."""
        GB_check_open(self)
        if not self.is_building:
            raise RuntimeError("Graph builder is not building.")
        cdef cydriver.CUstream c_stream = as_cu(self._h_stream)
        cdef cydriver.CUgraph c_graph
        with nogil:
            HANDLE_RETURN(cydriver.cuStreamEndCapture(c_stream, &c_graph))

        # TODO: Resolving https://github.com/NVIDIA/cuda-python/issues/617 would allow us to
        #       resume the build process after the first call to end_building()
        self._state = CAPTURE_ENDED
        return self

    def complete(self, options: GraphCompleteOptions | None = None) -> Graph:
        """Completes the graph builder and returns the built :obj:`~graph.Graph` object.

        Parameters
        ----------
        options : :obj:`~graph.GraphCompleteOptions`, optional
            Customizable dataclass for the graph builder completion options.

        Returns
        -------
        graph : :obj:`~graph.Graph`
            The newly built graph.

        """
        GB_check_open(self)
        if self._state != CAPTURE_ENDED:
            raise RuntimeError("Graph has not finished building.")

        return _instantiate_graph(self, options)

    def debug_dot_print(self, path: str, options: GraphDebugPrintOptions | None = None) -> None:
        """Generates a DOT debug file for the graph builder.

        Parameters
        ----------
        path : str
            File path to use for writting debug DOT output
        options : :obj:`~graph.GraphDebugPrintOptions`, optional
            Customizable dataclass for the debug print options.

        """
        GB_check_open(self)
        if self._state != CAPTURE_ENDED:
            raise RuntimeError("Graph has not finished building.")
        cdef unsigned int c_flags = options._to_flags() if options else 0
        cdef cydriver.CUgraph c_graph = as_cu(self._h_graph)
        cdef bytes b_path = path.encode('utf-8')
        cdef const char* c_path = b_path
        with nogil:
            HANDLE_RETURN(cydriver.cuGraphDebugDotPrint(c_graph, c_path, c_flags))

    def split(self, count: int) -> tuple[GraphBuilder, ...]:
        """Splits the original graph builder into multiple graph builders.

        The new builders inherit work dependencies from the original builder.
        The original builder is reused for the split and is returned first in the tuple.

        Parameters
        ----------
        count : int
            The number of graph builders to split the graph builder into.

        Returns
        -------
        graph_builders : tuple[:obj:`~graph.GraphBuilder`, ...]
            A tuple of split graph builders. The first graph builder in the tuple
            is always the original graph builder.

        """
        if count < 2:
            raise ValueError(f"Invalid split count: expecting >= 2, got {count}")
        GB_check_open(self)
        if self._state != CAPTURING:
            raise RuntimeError("Graph builder must be building before it can be split.")

        event = self._stream.record()
        result = [self]
        for i in range(count - 1):
            stream = self._stream.device.create_stream()
            stream.wait(event)
            result.append(GB_init_forked(stream, self._h_graph))
        event.close()
        return tuple(result)

    @staticmethod
    def join(*graph_builders: GraphBuilder) -> GraphBuilder:
        """Joins multiple graph builders into a single graph builder.

        The returned builder inherits work dependencies from the provided builders.

        Parameters
        ----------
        *graph_builders : :obj:`~graph.GraphBuilder`
            The graph builders to join.

        Returns
        -------
        graph_builder : :obj:`~graph.GraphBuilder`
            The newly joined graph builder.

        """
        if any(not isinstance(builder, GraphBuilder) for builder in graph_builders):
            raise TypeError("All arguments must be GraphBuilder instances")
        if len(graph_builders) < 2:
            raise ValueError("Must join with at least two graph builders")

        # Discover the root builder others should join
        root_idx = 0
        for i, builder in enumerate(graph_builders):
            if not builder.is_join_required:
                root_idx = i
                break

        # Join all onto the root builder
        root_bdr = graph_builders[root_idx]
        for idx, builder in enumerate(graph_builders):
            if idx == root_idx:
                continue
            root_bdr.stream.wait(builder.stream)
            builder.close()

        return root_bdr

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

    def _get_conditional_context(self) -> driver.CUcontext:
        return self._stream.context.handle

    def create_condition(self, default_value: int | None = None) -> GraphCondition:
        """Create a condition variable for use with conditional nodes.

        The returned :class:`GraphCondition` object is passed to conditional-node
        builder methods (:meth:`if_then`, :meth:`if_else`, :meth:`while_loop`,
        :meth:`switch`). Its value is controlled at runtime by device code via
        ``cudaGraphSetConditional``.

        Parameters
        ----------
        default_value : int, optional
            The default value to assign to the condition. If None, no
            default is assigned.

        Returns
        -------
        GraphCondition
            A condition variable for controlling conditional execution.
        """
        GB_check_open(self)
        if cy_driver_version() < (12, 3, 0):
            raise RuntimeError(f"Driver version {'.'.join(map(str, cy_driver_version()))} does not support conditional handles")
        if cy_binding_version() < (12, 3, 0):
            raise RuntimeError(f"Binding version {'.'.join(map(str, cy_binding_version()))} does not support conditional handles")
        if default_value is not None:
            flags = driver.CU_GRAPH_COND_ASSIGN_DEFAULT
        else:
            default_value = 0
            flags = 0

        status, _, graph, *_, _ = handle_return(driver.cuStreamGetCaptureInfo(self._stream.handle))
        if status != driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_ACTIVE:
            raise RuntimeError("Cannot create a condition when graph is not being built")

        raw_handle = handle_return(
            driver.cuGraphConditionalHandleCreate(graph, self._get_conditional_context(), default_value, flags)
        )
        return GraphCondition._from_handle(<cydriver.CUgraphConditionalHandle><intptr_t>int(raw_handle))

    def if_then(self, condition: GraphCondition) -> GraphBuilder:
        """Adds an if condition branch and returns a new graph builder for it.

        The resulting if graph will only execute the branch if the
        condition evaluates to true at runtime.

        The new builder inherits work dependencies from the original builder.

        Parameters
        ----------
        condition : :class:`~graph.GraphCondition`
            The condition variable from :meth:`create_condition` controlling
            whether the branch executes.

        Returns
        -------
        graph_builder : :obj:`~graph.GraphBuilder`
            The newly created conditional graph builder.

        """
        GB_check_open(self)
        if cy_driver_version() < (12, 3, 0):
            raise RuntimeError(f"Driver version {'.'.join(map(str, cy_driver_version()))} does not support conditional if")
        if cy_binding_version() < (12, 3, 0):
            raise RuntimeError(f"Binding version {'.'.join(map(str, cy_binding_version()))} does not support conditional if")
        if not isinstance(condition, GraphCondition):
            raise TypeError(
                f"condition must be a GraphCondition object (from "
                f"GraphBuilder.create_condition()), got {type(condition).__name__}")
        node_params = driver.CUgraphNodeParams()
        node_params.type = driver.CUgraphNodeType.CU_GRAPH_NODE_TYPE_CONDITIONAL
        node_params.conditional.handle = condition.handle
        node_params.conditional.type = driver.CUgraphConditionalNodeType.CU_GRAPH_COND_TYPE_IF
        node_params.conditional.size = 1
        node_params.conditional.ctx = self._get_conditional_context()
        return GB_cond_with_params(self, node_params)[0]

    def if_else(self, condition: GraphCondition) -> tuple[GraphBuilder, GraphBuilder]:
        """Adds an if-else condition branch and returns new graph builders for both branches.

        The resulting if graph will execute the branch if the condition
        evaluates to true at runtime, otherwise the else branch will execute.

        The new builders inherit work dependencies from the original builder.

        Parameters
        ----------
        condition : :class:`~graph.GraphCondition`
            The condition variable from :meth:`create_condition` controlling
            which branch executes.

        Returns
        -------
        graph_builders : tuple[:obj:`~graph.GraphBuilder`, :obj:`~graph.GraphBuilder`]
            A tuple of two new graph builders, one for the if branch and one for the else branch.

        """
        GB_check_open(self)
        if cy_driver_version() < (12, 8, 0):
            raise RuntimeError(f"Driver version {'.'.join(map(str, cy_driver_version()))} does not support conditional if-else")
        if cy_binding_version() < (12, 8, 0):
            raise RuntimeError(f"Binding version {'.'.join(map(str, cy_binding_version()))} does not support conditional if-else")
        if not isinstance(condition, GraphCondition):
            raise TypeError(
                f"condition must be a GraphCondition object (from "
                f"GraphBuilder.create_condition()), got {type(condition).__name__}")
        node_params = driver.CUgraphNodeParams()
        node_params.type = driver.CUgraphNodeType.CU_GRAPH_NODE_TYPE_CONDITIONAL
        node_params.conditional.handle = condition.handle
        node_params.conditional.type = driver.CUgraphConditionalNodeType.CU_GRAPH_COND_TYPE_IF
        node_params.conditional.size = 2
        node_params.conditional.ctx = self._get_conditional_context()
        return GB_cond_with_params(self, node_params)

    def switch(self, condition: GraphCondition, count: int) -> tuple[GraphBuilder, ...]:
        """Adds a switch condition branch and returns new graph builders for all cases.

        The resulting switch graph will execute the branch whose case index
        matches the value of the condition at runtime. If no match is found, no
        branch will be executed.

        The new builders inherit work dependencies from the original builder.

        Parameters
        ----------
        condition : :class:`~graph.GraphCondition`
            The condition variable from :meth:`create_condition` selecting
            which case executes.
        count : int
            The number of cases to add to the switch conditional.

        Returns
        -------
        graph_builders : tuple[:obj:`~graph.GraphBuilder`, ...]
            A tuple of new graph builders, one for each branch.

        """
        GB_check_open(self)
        if cy_driver_version() < (12, 8, 0):
            raise RuntimeError(f"Driver version {'.'.join(map(str, cy_driver_version()))} does not support conditional switch")
        if cy_binding_version() < (12, 8, 0):
            raise RuntimeError(f"Binding version {'.'.join(map(str, cy_binding_version()))} does not support conditional switch")
        if not isinstance(condition, GraphCondition):
            raise TypeError(
                f"condition must be a GraphCondition object (from "
                f"GraphBuilder.create_condition()), got {type(condition).__name__}")
        node_params = driver.CUgraphNodeParams()
        node_params.type = driver.CUgraphNodeType.CU_GRAPH_NODE_TYPE_CONDITIONAL
        node_params.conditional.handle = condition.handle
        node_params.conditional.type = driver.CUgraphConditionalNodeType.CU_GRAPH_COND_TYPE_SWITCH
        node_params.conditional.size = count
        node_params.conditional.ctx = self._get_conditional_context()
        return GB_cond_with_params(self, node_params)

    def while_loop(self, condition: GraphCondition) -> GraphBuilder:
        """Adds a while loop and returns a new graph builder for it.

        The resulting while loop graph will execute the branch repeatedly at runtime
        until the condition evaluates to false.

        The new builder inherits work dependencies from the original builder.

        Parameters
        ----------
        condition : :class:`~graph.GraphCondition`
            The condition variable from :meth:`create_condition` controlling
            loop continuation.

        Returns
        -------
        graph_builder : :obj:`~graph.GraphBuilder`
            The newly created while loop graph builder.

        """
        GB_check_open(self)
        if cy_driver_version() < (12, 3, 0):
            raise RuntimeError(f"Driver version {'.'.join(map(str, cy_driver_version()))} does not support conditional while loop")
        if cy_binding_version() < (12, 3, 0):
            raise RuntimeError(f"Binding version {'.'.join(map(str, cy_binding_version()))} does not support conditional while loop")
        if not isinstance(condition, GraphCondition):
            raise TypeError(
                f"condition must be a GraphCondition object (from "
                f"GraphBuilder.create_condition()), got {type(condition).__name__}")
        node_params = driver.CUgraphNodeParams()
        node_params.type = driver.CUgraphNodeType.CU_GRAPH_NODE_TYPE_CONDITIONAL
        node_params.conditional.handle = condition.handle
        node_params.conditional.type = driver.CUgraphConditionalNodeType.CU_GRAPH_COND_TYPE_WHILE
        node_params.conditional.size = 1
        node_params.conditional.ctx = self._get_conditional_context()
        return GB_cond_with_params(self, node_params)[0]

    def embed(self, GraphBuilder child):
        """Embed a previously-built :obj:`~graph.GraphBuilder` as a child node.

        Parameters
        ----------
        child : :obj:`~graph.GraphBuilder`
            The child graph builder. Must have finished building.
        """
        GB_check_open(self)
        if child._state != CAPTURE_ENDED:
            raise ValueError("Child graph has not finished building.")

        if not self.is_building:
            raise ValueError("Parent graph is not being built.")

        stream_handle = self._stream.handle
        _, _, graph_out, *deps_info_out, num_dependencies_out = handle_return(
            driver.cuStreamGetCaptureInfo(stream_handle)
        )

        # See https://github.com/NVIDIA/cuda-python/pull/879#issuecomment-3211054159
        # for rationale
        dependencies_out = deps_info_out[0]
        new_node = handle_return(
            driver.cuGraphAddChildGraphNode(
                graph_out, dependencies_out,
                num_dependencies_out, as_py(child._h_graph)
            )
        )
        cdef cydriver.CUgraphNode c_new_node = (
            <cydriver.CUgraphNode><intptr_t>int(new_node)
        )
        cdef cydriver.CUgraph embedded_graph = NULL
        cdef cydriver.CUresult rollback_status
        cdef GraphHandle h_embedded
        try:
            with nogil:
                HANDLE_RETURN(cydriver.cuGraphChildGraphNodeGetGraph(
                    c_new_node, &embedded_graph))
            h_embedded = create_child_graph_handle(
                embedded_graph, self._h_graph, c_new_node)
            HANDLE_RETURN(graph_clone_attachments(
                h_embedded, child._h_graph))
        except:
            with nogil:
                rollback_status = cydriver.cuGraphDestroyNode(c_new_node)
            if rollback_status == cydriver.CUDA_SUCCESS:
                invalidate_child_graph_state(
                    self._h_graph, c_new_node)
            raise

        deps_info_update = [[new_node]] + [None] * (len(deps_info_out) - 1)
        handle_return(
            driver.cuStreamUpdateCaptureDependencies(
                stream_handle,
                *deps_info_update,  # dependencies, edgeData
                1,
                driver.CUstreamUpdateCaptureDependencies_flags.CU_STREAM_SET_CAPTURE_DEPENDENCIES,
            )
        )

    def callback(self, fn, *, user_data=None) -> None:
        """Add a host callback to the graph during stream capture.

        The callback runs on the host CPU when the graph reaches this point
        in execution. Two modes are supported:

        - **Python callable**: Pass any callable. The GIL is acquired
          automatically. The callable must take no arguments; use closures
          or ``functools.partial`` to bind state.
        - **ctypes function pointer**: The function receives a single
          ``void*`` argument (the ``user_data``), and the caller must keep
          the ctypes wrapper alive for the lifetime of the graph. Its
          declared prototype must match the driver's ``CUhostFn``
          (``void (*)(void*)``): ``ctypes.CFUNCTYPE(None, ctypes.c_void_p)``,
          or ``ctypes.WINFUNCTYPE(None, ctypes.c_void_p)`` on Windows.

        .. warning::

            Callbacks must not call CUDA API functions. Doing so may
            deadlock or corrupt driver state.

            Use caution when a Python callback retains an object that owns a
            graph. Any reference cycle involving the callback and a graph that
            retains it cannot be broken by Python's cyclic garbage collector.
            Use a weak reference to break such cycles.

        Parameters
        ----------
        fn : callable or ctypes function pointer
            The callback function.
        user_data : int or bytes-like, optional
            Only for ctypes function pointers. If ``int``, passed as a raw
            pointer (caller manages lifetime). If bytes-like, the data is
            copied and its lifetime is tied to the graph.

        Raises
        ------
        TypeError
            If ``fn`` is a ctypes function pointer whose declared prototype
            does not match ``CUhostFn``.
        ValueError
            If ``user_data`` is given for a Python callable.
        """
        GB_callback(self, fn, user_data, False)


cdef inline void GB_callback(
        GraphBuilder gb, object fn, object user_data,
        bint fail_tail_discovery_for_testing) except *:
    GB_check_open(gb)
    cdef Stream stream = gb._stream
    cdef cydriver.CUstream c_stream = as_cu(stream._h_stream)
    cdef cydriver.CUstreamCaptureStatus capture_status

    with nogil:
        _get_capture_info(c_stream, &capture_status, NULL)

    if capture_status != cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE:
        raise RuntimeError("Cannot add callback when graph is not being built")

    cdef cydriver.CUhostFn c_fn
    cdef void* c_user_data = NULL
    cdef OpaqueHandle fn_owner, data_owner
    cdef PreparedAttachment prepared
    _resolve_host_callback(
        fn, user_data, &c_fn, &c_user_data, &fn_owner, &data_owner)
    HANDLE_RETURN(graph_prepare_attachment(
        gb._h_graph, fn_owner, data_owner, &prepared))

    with nogil:
        HANDLE_RETURN(cydriver.cuLaunchHostFunc(c_stream, c_fn, c_user_data))

    # Capturing the host function added a node to the graph; it is now the
    # stream's sole capture dependency. Attach the callback's owners to it.
    cdef cydriver.CUgraphNode host_node
    cdef cydriver.CUresult commit_status
    try:
        if fail_tail_discovery_for_testing:
            raise RuntimeError("forced capture tail discovery failure")
        host_node = _capture_tail_node(c_stream)
    except:
        # CUDA added the callback, but its node cannot be identified.
        # Retain its owners anonymously to prevent dangling pointers.
        commit_status = graph_commit_attachment(prepared, NULL)
        HANDLE_RETURN(commit_status)
        raise
    HANDLE_RETURN(graph_commit_attachment(prepared, host_node))


def _capture_callback_with_tail_failure_for_testing(
        GraphBuilder gb, fn, *, user_data=None):
    """Exercise anonymous attachment retention after node discovery fails."""
    GB_callback(gb, fn, user_data, True)


cdef inline int GB_check_open(GraphBuilder gb) except -1:
    """Reject operations on a builder that has been closed.

    A CLOSED builder has reset its stream and graph handles, so any method
    that dereferences them would read a null handle (or, for the cached
    Stream, a None typed as cdef Stream). Guarding here yields a clear error
    instead.
    """
    if gb._state == CLOSED:
        raise RuntimeError("Graph builder has been closed.")
    return 0


cdef inline int GB_end_capture_if_needed(GraphBuilder gb, bint check_status) except -1 nogil:
    """End an in-progress capture if this builder owns it.

    Only a CAPTURING PRIMARY or CONDITIONAL_BODY builder owns the live
    capture. A FORKED builder must not call cuStreamEndCapture: the driver
    requires forked streams to be joined first.

    check_status=True checks the driver return (close()); False ignores it
    (__dealloc__).
    """
    cdef cydriver.CUgraph c_graph
    cdef cydriver.CUresult err
    cdef cydriver.CUstream c_stream
    if gb._h_stream and gb._state == CAPTURING and gb._kind != FORKED:
        c_stream = as_cu(gb._h_stream)
        with nogil:
            err = cydriver.cuStreamEndCapture(c_stream, &c_graph)
            if check_status:
                HANDLE_RETURN(err)
    return 0


cdef inline GraphBuilder GB_init_forked(Stream stream, GraphHandle h_primary_graph):
    cdef GraphBuilder gb = GraphBuilder.__new__(GraphBuilder)
    # A FORKED builder captures into the primary's CUgraph. It holds the
    # primary's GraphHandle so conditional bodies share its graph hierarchy.
    gb._h_graph = h_primary_graph
    gb._h_stream = stream._h_stream
    gb._kind = FORKED
    gb._state = CAPTURING
    gb._stream = stream
    return gb


cdef inline GraphBuilder GB_init_conditional(
        Stream stream, cydriver.CUgraph cond_graph,
        GraphBuilder parent, cydriver.CUgraphNode owner_node):
    cdef GraphBuilder gb = GraphBuilder.__new__(GraphBuilder)
    gb._h_graph = create_child_graph_handle(
        cond_graph, parent._h_graph, owner_node)
    gb._h_stream = stream._h_stream
    gb._kind = CONDITIONAL_BODY
    gb._state = CAPTURE_NOT_STARTED
    gb._stream = stream
    return gb


cdef inline int _get_capture_info(
        cydriver.CUstream stream,
        cydriver.CUstreamCaptureStatus* status,
        cydriver.CUgraph* graph) except?-1 nogil:
    """Thin wrapper around ``cuStreamGetCaptureInfo`` that papers over the
    CUDA 12 vs 13 signature change.

    ``status`` must be non-NULL: the driver rejects ``captureStatus_out=NULL``
    with ``CUDA_ERROR_INVALID_VALUE``. ``graph`` may be NULL when the caller
    does not need the graph handle.
    """
    IF CUDA_CORE_BUILD_MAJOR >= 13:
        return HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(
            stream, status, NULL, graph, NULL, NULL, NULL))
    ELSE:
        return HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(
            stream, status, NULL, graph, NULL, NULL))


cdef inline cydriver.CUgraphNode _capture_tail_node(cydriver.CUstream stream) except *:
    """Return the node a freshly-captured single-node operation left as the
    stream's sole capture dependency (e.g. the host node added by
    ``cuLaunchHostFunc``). The driver advances the stream's dependency set to
    the new node, so the next captured op would depend on it.
    """
    cdef cydriver.CUstreamCaptureStatus status
    cdef const cydriver.CUgraphNode* deps = NULL
    cdef size_t num_deps = 0
    with nogil:
        IF CUDA_CORE_BUILD_MAJOR >= 13:
            HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(
                stream, &status, NULL, NULL, &deps, NULL, &num_deps))
        ELSE:
            HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(
                stream, &status, NULL, NULL, &deps, &num_deps))
    if num_deps != 1:
        raise RuntimeError(
            f"expected exactly one capture dependency after a host callback, got {num_deps}")
    return <cydriver.CUgraphNode>deps[0]


cdef inline tuple GB_cond_with_params(GraphBuilder gb, node_params):
    status, _, graph, *deps_info, num_dependencies = handle_return(
        driver.cuStreamGetCaptureInfo(gb._stream.handle)
    )
    if status != driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_ACTIVE:
        raise RuntimeError("Cannot add conditional node when not actively capturing")

    new_node = handle_return(
        driver.cuGraphAddNode(graph, *deps_info, num_dependencies, node_params))
    deps_info_update = [[new_node]] + [None] * (len(deps_info) - 1)

    handle_return(
        driver.cuStreamUpdateCaptureDependencies(
            gb._stream.handle,
            *deps_info_update,  # dependencies, edgeData
            1,  # numDependencies
            driver.CUstreamUpdateCaptureDependencies_flags.CU_STREAM_SET_CAPTURE_DEPENDENCIES,
        )
    )

    return tuple(
        GB_init_conditional(
            gb._stream.device.create_stream(),
            <cydriver.CUgraph><intptr_t>int(node_params.conditional.phGraph_out[i]),
            gb,
            <cydriver.CUgraphNode><intptr_t>int(new_node),
        )
        for i in range(node_params.conditional.size)
    )


cdef class Graph:
    """An executable graph.

    A graph groups a set of CUDA kernels and other CUDA operations together and executes
    them with a specified dependency tree. It speeds up the workflow by combining the
    driver activities associated with CUDA kernel launches and CUDA API calls.

    Graphs must be built using a :obj:`~graph.GraphBuilder` object.

    """

    def __init__(self):
        raise RuntimeError("directly constructing a Graph instance is not supported")

    @staticmethod
    cdef Graph _init(GraphExecHandle h_graph_exec):
        cdef Graph self = Graph.__new__(Graph)
        self._h_graph_exec = h_graph_exec
        return self

    def close(self) -> None:
        """Destroy the graph."""
        self._h_graph_exec.reset()
        retry_deferred_cleanup()

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

        .. caution::

            This handle is a Python object. To get the memory address of the underlying C
            handle, call ``int()`` on the returned object.

        """
        return as_py(self._h_graph_exec)

    def __getitem__(self, node: GraphNode) -> ExecutableGraphNode:
        """Return a view for updating *node* in this executable graph.

        *node* is a definition node from the graph used to instantiate this
        executable. Call ``update()`` on the returned view to replace that
        node's parameters for future launches. Kernel, memcpy, and memset
        views also support enabling and disabling the node.
        """
        return create_executable_node_view(
            self._h_graph_exec, node)

    def update(self, source: "GraphBuilder | GraphDefinition") -> None:
        """Update the graph using a new graph definition.

        The topology of the provided source must be identical to this graph.

        Parameters
        ----------
        source : :obj:`~graph.GraphBuilder` or :obj:`~graph.GraphDefinition`
            The graph definition to update from. A GraphBuilder must have
            finished building.

        """
        cdef GraphHandle h_source

        if isinstance(source, GraphBuilder):
            if (<GraphBuilder>source)._state == CLOSED:
                raise ValueError("Source graph builder has been closed.")
            if (<GraphBuilder>source)._state != CAPTURE_ENDED:
                raise ValueError("Graph has not finished building.")
            h_source = (<GraphBuilder>source)._h_graph
        elif isinstance(source, GraphDefinition):
            h_source = (<GraphDefinition>source)._h_graph
        else:
            raise TypeError(
                f"expected GraphBuilder or GraphDefinition, got {type(source).__name__}")

        cdef cydriver.CUgraphExecUpdateResultInfo result_info
        cdef cydriver.CUresult err = graph_exec_update(
            self._h_graph_exec, h_source, &result_info)
        if err == cydriver.CUresult.CUDA_ERROR_GRAPH_EXEC_UPDATE_FAILURE:
            reason = driver.CUgraphExecUpdateResult(result_info.result)
            msg = f"Graph update failed: {reason.__doc__.strip()} ({reason.name})"
            raise CUDAError(msg)
        HANDLE_RETURN(err)

    def upload(self, stream: Stream) -> None:
        """Uploads the graph in a stream.

        Parameters
        ----------
        stream : :obj:`~_stream.Stream`
            The stream in which to upload the graph

        """
        cdef cydriver.CUgraphExec c_exec = as_cu(self._h_graph_exec)
        cdef cydriver.CUstream c_stream = <cydriver.CUstream><intptr_t>int(stream.handle)
        with nogil:
            HANDLE_RETURN(cydriver.cuGraphUpload(c_exec, c_stream))

    def launch(self, stream: Stream) -> None:
        """Launches the graph in a stream.

        Parameters
        ----------
        stream : :obj:`~_stream.Stream`
            The stream in which to launch the graph.

        """
        cdef cydriver.CUgraphExec c_exec = as_cu(self._h_graph_exec)
        cdef cydriver.CUstream c_stream = <cydriver.CUstream><intptr_t>int(stream.handle)
        with nogil:
            HANDLE_RETURN(cydriver.cuGraphLaunch(c_exec, c_stream))
