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

from __future__ import annotations

from collections.abc import Callable, Iterable, Iterator, Set
from dataclasses import dataclass
from typing import Any

from cuda.core._device import Device
from cuda.core._memory._device_memory_resource import DeviceMemoryResource


@dataclass(frozen=True)
class PeerAccessPlan:
    """Normalized peer-access target state and the driver updates it requires."""
    target_ids: tuple[int, ...]
    to_add: tuple[int, ...]
    to_remove: tuple[int, ...]

class PeerAccessibleBySetProxy:
    """Live driver-backed view of the peer devices granted access to a memory pool.

    Reads (``__contains__``, ``__iter__``, ``len(...)``) call ``cuMemPoolGetAccess``;
    writes (``add``, ``discard``, and bulk ops) call ``cuMemPoolSetAccess``. There
    is no in-memory mirror, so the view always reflects the current driver state
    and stays consistent across multiple wrappers around the same pool.

    Iteration yields :class:`~cuda.core.Device` objects. ``add``, ``discard``, and
    ``__contains__`` accept either a :class:`~cuda.core.Device` or a device-ordinal
    ``int``; the owner device is silently ignored when supplied.

    All bulk operations (``update``, ``|=``, ``&=``, ``-=``, ``^=``, ``clear``)
    issue exactly one ``cuMemPoolSetAccess`` call. This matters: peer-access
    transitions can take seconds per pool because every existing memory mapping
    is updated, so coalescing into a single driver call lets the toolkit handle
    the mappings in parallel.
    """
    __slots__ = ('_mr',)

    def __init__(self, mr: DeviceMemoryResource) -> None:
        ...

    @classmethod
    def _from_iterable(cls, it: Iterable[Device]) -> set[Device]:
        ...

    def __contains__(self, value: object) -> bool:
        ...

    def __iter__(self) -> Iterator[Device]:
        ...

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

    def add(self, value: Device | int) -> None:
        """Grant peer access from ``value`` to allocations in this pool."""

    def discard(self, value: Device | int) -> None:
        """Revoke peer access from ``value`` to allocations in this pool."""

    def clear(self) -> None:
        """Revoke all peer access in a single driver call."""

    def update(self, *others: Iterable[Device | int]) -> None:
        """Grant peer access to every device in ``others`` in one driver call."""

    def difference_update(self, *others: Iterable[Device | int]) -> None:
        """Revoke peer access for every device in ``others`` in one driver call."""

    def intersection_update(self, *others: Iterable[Device | int]) -> None:
        """Restrict peer access to the intersection in a single driver call."""

    def symmetric_difference_update(self, other: Iterable[Device | int]) -> None:
        """Toggle peer access for every device in ``other`` in one driver call."""

    def __ior__(self, other: Set[Any]) -> PeerAccessibleBySetProxy:
        ...

    def __iand__(self, other: Set[Any]) -> PeerAccessibleBySetProxy:
        ...

    def __isub__(self, other: Set[Any]) -> PeerAccessibleBySetProxy:
        ...

    def __ixor__(self, other: Set[Any]) -> PeerAccessibleBySetProxy:
        ...

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

    def _apply(self, additions, removals) -> None:
        """Compute the diff and issue a single ``cuMemPoolSetAccess``.

        ``additions`` and ``removals`` are user-supplied (``Device | int``);
        only the owner device is filtered out. Adds are validated through
        :meth:`Device.can_access_peer` via :func:`plan_peer_access_update`;
        removals bypass that check (revoking is always permitted).
        """

def replace_peer_accessible_by(mr: DeviceMemoryResource, devices: object) -> None:
    """Replace the full peer-access set in a single batched driver call.

    Backs the ``mr.peer_accessible_by = [...]`` setter. Uses the same planner
    as the proxy's bulk ops; the only difference is that adds and removes are
    derived from the symmetric difference between current driver state and the
    requested target set.
    """

def normalize_peer_access_targets(owner_device_id: int, requested_devices: Iterable[object], *, resolve_device_id: Callable[[object], int]) -> tuple[int, ...]:
    """Return sorted, unique peer device IDs, excluding the owner device."""

def plan_peer_access_update(owner_device_id: int, current_peer_ids: Iterable[int], requested_devices: Iterable[object], *, resolve_device_id: Callable[[object], int], can_access_peer: Callable[[int], bool]) -> PeerAccessPlan:
    """Compute the peer-access target state and add/remove deltas."""

def _resolve_peer_device_id(value: Device | int | None) -> int:
    """Coerce ``Device | int`` into a device-ordinal int."""

def _set_pool_access(mr: object, to_add: tuple[int, ...], to_remove: tuple[int, ...]) -> None:
    """Issue one ``cuMemPoolSetAccess`` for the given add/remove deltas.

    The thin Python-callable layer that wraps the actual driver call: building
    the ``CUmemAccessDesc`` array and invoking ``cuMemPoolSetAccess`` happens
    in here. Tests monkeypatch this on the module to spy on real driver work
    without intercepting earlier no-op paths.

    Preconditions: ``len(to_add) + len(to_remove) > 0`` (the caller is
    responsible for skipping empty diffs).
    """

def _apply_peer_access_diff(mr: DeviceMemoryResource, to_add: Iterable[int], to_remove: Iterable[int]) -> None:
    """Apply a peer-access diff in at most one driver call.

    Every write path on :class:`PeerAccessibleBySetProxy` and the
    ``peer_accessible_by`` setter routes through this function. Empty diffs
    short-circuit here so the driver-level helper :func:`_set_pool_access` is
    only invoked when there is actual work for ``cuMemPoolSetAccess`` to do.
    """