from __future__ import annotations

from typing import Tuple, TypeVar

# backwards compatibility. Luckily circular references are fine in type stubs
from ray._raylet import ObjectRef

ObjectID = ObjectRef

# implementations are in unique_ids.pxi
def check_id(b: bytes, size: int = ...) -> None: ...

_BID = TypeVar("_BID", bound=BaseID)
class BaseID:

    @classmethod
    def from_binary(cls: type[_BID], id_bytes: bytes) -> _BID: ...

    @classmethod
    def from_hex(cls: type[_BID], hex_id: str | bytes) -> _BID: ...

    def binary(self) -> bytes: ...

    @classmethod
    def size(cls) -> int: ...

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

    def is_nil(self) -> bool: ...

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

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

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

    def __bytes__(self) -> bytes: ...

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

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

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

    def __reduce__(self: _BID) -> Tuple[type[_BID], Tuple[bytes]]: ...

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


_UID = TypeVar("_UID", bound=UniqueID)
class UniqueID(BaseID):

    def __init__(self, id: bytes) -> None: ...

    @classmethod
    def nil(cls: type[_UID]) -> _UID: ...

    @classmethod
    def from_random(cls: type[_UID]) -> _UID: ...


_TID = TypeVar("_TID", bound=TaskID)
class TaskID(BaseID):

    def __init__(self, id: bytes) -> None: ...

    def actor_id(self) -> ActorID: ...

    def job_id(self) -> JobID: ...

    @classmethod
    def nil(cls: type[_TID]) -> _TID: ...

    @classmethod
    def for_fake_task(cls: type[_TID], job_id: JobID) -> _TID: ...

    @classmethod
    def for_driver_task(cls: type[_TID], job_id: JobID) -> _TID: ...

    @classmethod
    def for_actor_creation_task(cls: type[_TID], actor_id: ActorID) -> _TID: ...

    @classmethod
    def for_actor_task(cls: type[_TID], job_id: JobID, parent_task_id: TaskID,
                       parent_task_counter: int, actor_id: ActorID) -> _TID: ...

    @classmethod
    def for_normal_task(cls: type[_TID], job_id: JobID, parent_task_id: TaskID, parent_task_counter: int) -> _TID: ...


class NodeID(UniqueID): ...

_JID = TypeVar("_JID", bound=JobID)
class JobID(BaseID):

    def __init__(self, id: bytes) -> None: ...

    @classmethod
    def from_int(cls: type[_JID], value: int) -> _JID: ...

    @classmethod
    def nil(cls: type[_JID]) -> _JID: ...

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


class WorkerID(UniqueID): ...

_AID = TypeVar("_AID", bound=ActorID)
class ActorID(BaseID):

    def __init__(self, id: bytes) -> None: ...

    @classmethod
    def of(cls: type[_AID], job_id: JobID, parent_task_id: TaskID, parent_task_counter: int) -> _AID: ...

    @classmethod
    def nil(cls: type[_AID]) -> _AID: ...

    @classmethod
    def from_random(cls: type[_AID]) -> _AID: ...

    def _set_id(self, id: bytes) -> None: ...

    @property
    def job_id(self) -> JobID: ...


class FunctionID(UniqueID): ...
class ActorClassID(UniqueID): ...
class ClusterID(UniqueID): ...


_PGID = TypeVar("_PGID", bound=PlacementGroupID)
class PlacementGroupID(BaseID):

    def __init__(self, id: bytes) -> None: ...

    @classmethod
    def from_random(cls: type[_PGID]) -> _PGID: ...

    @classmethod
    def of(cls: type[_PGID], job_id: JobID) -> _PGID: ...

    @classmethod
    def nil(cls: type[_PGID]) -> _PGID: ...