# This file was generated by stubgen-pyx v0.2.6 from cuda_core/cuda/core/graph/_graph_builder.pyx

from __future__ import annotations

from dataclasses import dataclass

from cuda.core._stream import Stream
from cuda.core._utils.cuda_utils import driver
from cuda.core.graph._graph_definition import GraphCondition, GraphDefinition

_BuilderKind = int
_CaptureState = int

@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)."""

@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

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):
        ...

    def __dealloc__(self):
        ...

    @staticmethod
    def _init(stream: Stream):
        ...

    def close(self):
        """Destroy the graph builder."""

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

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

    @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.
        """

    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.

        """

    @property
    def is_building(self) -> bool:
        """Returns True if the graph builder is currently building."""

    def end_building(self) -> GraphBuilder:
        """Ends the building process."""

    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.

        """

    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.

        """

    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.

        """

    @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.

        """

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

    def _get_conditional_context(self) -> driver.CUcontext:
        ...

    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.
        """

    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.

        """

    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.

        """

    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.

        """

    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.

        """

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

        Parameters
        ----------
        child : :obj:`~graph.GraphBuilder`
            The child graph builder. Must have finished building.
        """

    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**: Pass a ``ctypes.CFUNCTYPE`` instance.
          The function receives a single ``void*`` argument (the
          ``user_data``). The caller must keep the ctypes wrapper alive
          for the lifetime of the graph.

        .. warning::

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

        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.
        """

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):
        ...

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

    @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.

        """

    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.

        """

    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

        """

    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.

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

def _instantiate_graph(h_graph, options: GraphCompleteOptions | None=None) -> Graph:
    ...