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

"""GraphNode base class — factory, properties, and builder methods."""
from __future__ import annotations

import weakref
from collections.abc import Iterable

from cuda.core._device import Device
from cuda.core._event import Event
from cuda.core._launch_config import LaunchConfig
from cuda.core._memory._buffer import Buffer
from cuda.core._module import Kernel
from cuda.core._utils.cuda_utils import driver
from cuda.core.graph._adjacency_set_proxy import AdjacencySetProxy
from cuda.core.graph._graph_definition import GraphCondition, GraphDefinition
from cuda.core.graph._subclasses import (AllocNode, ChildGraphNode, EmptyNode,
                                         EventRecordNode, EventWaitNode,
                                         FreeNode, IfElseNode, IfNode,
                                         KernelNode, MemcpyNode, MemsetNode,
                                         SwitchNode, WhileNode)
from cuda.core.typing import GraphMemoryType


class GraphNode:
    """A node in a graph definition.

    Nodes are created by calling builder methods on GraphDefinition (for
    entry-point nodes with no dependencies) or on other Nodes (for
    nodes that depend on a predecessor).
    """

    def __repr__(self) -> str:
        ...

    def __eq__(self, other: object) -> bool:
        ...

    def __hash__(self) -> int:
        ...

    @property
    def type(self) -> driver.CUgraphNodeType | None:
        """Return the CUDA graph node type.

        Returns
        -------
        CUgraphNodeType or None
            The node type enum value, or None for the entry node.
        """

    @property
    def graph(self) -> GraphDefinition:
        """Return the GraphDefinition this node belongs to."""

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

        Returns None for the entry node.
        """

    @property
    def is_valid(self) -> bool:
        """Whether this node is valid (not destroyed).

        Returns ``False`` after :meth:`destroy` has been called.
        """

    def destroy(self) -> None:
        """Destroy this node and remove all its edges from the parent graph.

        After this call, :attr:`is_valid` returns ``False`` and the node
        cannot be re-added to any graph.  Safe to call on an
        already-destroyed node (no-op).
        """

    @property
    def pred(self) -> AdjacencySetProxy:
        """A mutable set-like view of this node's predecessors."""

    @pred.setter
    def pred(self, value: Iterable[GraphNode]) -> None:
        ...

    @property
    def succ(self) -> AdjacencySetProxy:
        """A mutable set-like view of this node's successors."""

    @succ.setter
    def succ(self, value: Iterable[GraphNode]) -> None:
        ...

    def launch(self, config: LaunchConfig, kernel: Kernel, *args) -> KernelNode:
        """Add a kernel launch node depending on this node.

        Parameters
        ----------
        config : LaunchConfig
            Launch configuration (grid, block, shared memory, etc.)
        kernel : Kernel
            The kernel to launch.
        *args
            Kernel arguments.

        Returns
        -------
        KernelNode
            A new KernelNode representing the kernel launch.
        """

    def join(self, *nodes: GraphNode) -> EmptyNode:
        """Create an empty node that depends on this node and all given nodes.

        This is used to synchronize multiple branches of execution.

        Parameters
        ----------
        *nodes : GraphNode
            Additional nodes to depend on.

        Returns
        -------
        EmptyNode
            A new EmptyNode that depends on all input nodes.
        """

    def allocate(self, size: int, *, device: Device | int | None=None, memory_type: GraphMemoryType=..., peer_access: list[Device | int] | None=None) -> AllocNode:
        """Add a memory allocation node depending on this node.

        Parameters
        ----------
        size : int
            Number of bytes to allocate.
        device : int or Device, optional
            The device on which to allocate memory. If None (default),
            uses the current CUDA context's device.
        memory_type : GraphMemoryType or str, optional
            Type of memory to allocate. One of:

            - ``GraphMemoryType.DEVICE`` (default): Pinned device memory,
              optimal for GPU kernels.
            - ``GraphMemoryType.HOST``: Pinned host memory, accessible from
              both host and device. Useful for graphs containing host
              callback nodes. Note: may not be supported on all
              systems/drivers.
            - ``GraphMemoryType.MANAGED``: Managed/unified memory that
              automatically migrates between host and device. Useful for
              mixed host/device access patterns.

        peer_access : list of int or Device, optional
            List of devices that should have read-write access to the
            allocated memory. If None (default), only the allocating
            device has access.

        Returns
        -------
        AllocNode
            A new AllocNode representing the allocation. Access the allocated
            device pointer via the dptr property.

        Notes
        -----
        IPC (inter-process communication) is not supported for graph
        memory allocation nodes per CUDA documentation.
        """

    def deallocate(self, dptr: int) -> FreeNode:
        """Add a memory free node depending on this node.

        Parameters
        ----------
        dptr : int
            Device pointer to free (typically from AllocNode.dptr).

        Returns
        -------
        FreeNode
            A new FreeNode representing the free operation.
        """

    def memset(self, dst: Buffer | int, value, width: int, height: int=1, pitch: int=0, *, dst_owner=None) -> MemsetNode:
        """Add a memset node depending on this node.

        Parameters
        ----------
        dst : Buffer or int
            Destination. When ``dst`` is a :class:`Buffer`, the underlying
            allocation is retained for the graph's lifetime. A raw pointer
            (``int``) is used as-is; the caller must keep the underlying memory
            alive, or supply ``dst_owner`` to have the graph retain it.
        value : int or buffer-protocol object
            Fill value. int for 1-byte fill (range [0, 256)),
            or buffer-protocol object of 1, 2, or 4 bytes.
        width : int
            Width of the row in elements.
        height : int, optional
            Number of rows (default 1).
        pitch : int, optional
            Pitch of destination in bytes (default 0, unused if height is 1).
        dst_owner : object, optional
            Object retained for the graph's lifetime when ``dst`` is a raw
            pointer. A :class:`Buffer` owner retains its underlying allocation,
            not the wrapper. Must not be passed when ``dst`` is a :class:`Buffer`.

        Returns
        -------
        MemsetNode
            A new MemsetNode representing the memset operation.

        Raises
        ------
        ValueError
            If ``dst_owner`` is given together with a :class:`Buffer` ``dst``.
        """

    def memcpy(self, dst: Buffer | int, src: Buffer | int, size: int, *, dst_owner=None, src_owner=None) -> MemcpyNode:
        """Add a memcpy node depending on this node.

        Copies ``size`` bytes from ``src`` to ``dst``. Memory types are
        auto-detected via the driver, so both device and pinned host
        pointers are supported.

        Parameters
        ----------
        dst : Buffer or int
            Destination (device or pinned host). When a :class:`Buffer` is given,
            the underlying allocation is retained for the graph's lifetime. A raw
            pointer (``int``) is used as-is; the caller must keep the underlying
            memory alive, or supply ``dst_owner`` to have the graph retain it.
        src : Buffer or int
            Source (device or pinned host). Same retention rules as ``dst``;
            use ``src_owner`` for a raw pointer.
        size : int
            Number of bytes to copy.
        dst_owner : object, optional
            Object retained for the graph's lifetime when ``dst`` is a raw
            pointer. A :class:`Buffer` owner retains its underlying allocation.
            Must not be passed when ``dst`` is a :class:`Buffer`.
        src_owner : object, optional
            Object retained for the graph's lifetime when ``src`` is a raw
            pointer. A :class:`Buffer` owner retains its underlying allocation.
            Must not be passed when ``src`` is a :class:`Buffer`.

        Returns
        -------
        MemcpyNode
            A new MemcpyNode representing the copy operation.

        Raises
        ------
        ValueError
            If ``dst_owner`` or ``src_owner`` is given together with a
            :class:`Buffer` ``dst`` or ``src`` respectively.
        """

    def embed(self, child: GraphDefinition) -> ChildGraphNode:
        """Add a child graph node depending on this node.

        Embeds a clone of the given graph definition as a sub-graph node.
        The child graph must not contain allocation, free, or conditional
        nodes.

        Parameters
        ----------
        child : GraphDefinition
            The graph definition to embed (will be cloned).

        Returns
        -------
        ChildGraphNode
            A new ChildGraphNode representing the embedded sub-graph.
        """

    def record(self, event: Event) -> EventRecordNode:
        """Add an event record node depending on this node.

        Parameters
        ----------
        event : Event
            The event to record.

        Returns
        -------
        EventRecordNode
            A new EventRecordNode representing the event record operation.
        """

    def wait(self, event: Event) -> EventWaitNode:
        """Add an event wait node depending on this node.

        Parameters
        ----------
        event : Event
            The event to wait for.

        Returns
        -------
        EventWaitNode
            A new EventWaitNode representing the event wait operation.
        """

    def callback(self, fn, *, user_data=None) -> object:
        """Add a host callback node depending on this node.

        The callback runs on the host CPU when the graph reaches this node.
        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.

        Returns
        -------
        HostCallbackNode
            A new HostCallbackNode representing the callback.
        """

    def if_then(self, condition: GraphCondition) -> IfNode:
        """Add an if-conditional node depending on this node.

        The body graph executes only when the condition evaluates to
        a non-zero value at runtime.

        Parameters
        ----------
        condition : GraphCondition
            GraphCondition from :meth:`GraphDefinition.create_condition`.

        Returns
        -------
        IfNode
            A new IfNode with one branch accessible via ``.then``.
        """

    def if_else(self, condition: GraphCondition) -> IfElseNode:
        """Add an if-else conditional node depending on this node.

        Two body graphs: the first executes when the condition is
        non-zero, the second when it is zero.

        Parameters
        ----------
        condition : GraphCondition
            GraphCondition from :meth:`GraphDefinition.create_condition`.

        Returns
        -------
        IfElseNode
            A new IfElseNode with branches accessible via
            ``.then`` and ``.else_``.
        """

    def while_loop(self, condition: GraphCondition) -> WhileNode:
        """Add a while-loop conditional node depending on this node.

        The body graph executes repeatedly while the condition
        evaluates to a non-zero value.

        Parameters
        ----------
        condition : GraphCondition
            GraphCondition from :meth:`GraphDefinition.create_condition`.

        Returns
        -------
        WhileNode
            A new WhileNode with body accessible via ``.body``.
        """

    def switch(self, condition: GraphCondition, count: int) -> SwitchNode:
        """Add a switch conditional node depending on this node.

        The condition value selects which branch to execute. If the
        value is out of range, no branch executes.

        Parameters
        ----------
        condition : GraphCondition
            GraphCondition from :meth:`GraphDefinition.create_condition`.
        count : int
            Number of switch cases (branches).

        Returns
        -------
        SwitchNode
            A new SwitchNode with branches accessible via ``.branches``.
        """
__all__ = ['GraphNode']
_node_registry: weakref.WeakValueDictionary[int, GraphNode] = weakref.WeakValueDictionary()