"""Type stubs for spork.runtime.pds C extension."""

from collections.abc import Iterator
from typing import Any, Generic, TypeVar

T = TypeVar("T")
K = TypeVar("K")
V = TypeVar("V")

# =============================================================================
# Cons - Persistent linked list cell
# =============================================================================

class Cons(Generic[T]):
    first: T
    rest: Cons[T] | None

    def __init__(self, first: T, rest: Cons[T] | None = None) -> None: ...
    def __iter__(self) -> Iterator[T]: ...
    def __len__(self) -> int: ...
    def __hash__(self) -> int: ...
    def __eq__(self, other: object) -> bool: ...
    def __ne__(self, other: object) -> bool: ...
    def __repr__(self) -> str: ...
    def __class_getitem__(cls, item: Any) -> Any: ...
    def conj(self, val: T) -> Cons[T]: ...

# =============================================================================
# Vector - Persistent vector (bit-partitioned trie)
# =============================================================================

class Vector(Generic[T]):
    def __init__(self, *args: T) -> None: ...
    def __iter__(self) -> Iterator[T]: ...
    def __len__(self) -> int: ...
    def __getitem__(self, index: int) -> T: ...
    def __hash__(self) -> int: ...
    def __eq__(self, other: object) -> bool: ...
    def __ne__(self, other: object) -> bool: ...
    def __repr__(self) -> str: ...
    def __contains__(self, item: object) -> bool: ...
    def __add__(self, other: Vector[T]) -> Vector[T]: ...
    def __class_getitem__(cls, item: Any) -> Any: ...
    def nth(self, index: int, default: T = ...) -> T: ...
    def conj(self, val: T) -> Vector[T]: ...
    def assoc(self, index: int, val: T) -> Vector[T]: ...
    def pop(self) -> Vector[T]: ...
    def transient(self) -> TransientVector[T]: ...
    def to_seq(self) -> Cons[T]: ...
    def copy(self) -> Vector[T]: ...
    def index(self, value: T, start: int = 0, stop: int = ...) -> int: ...
    def count(self, value: T) -> int: ...

class TransientVector(Generic[T]):
    def __len__(self) -> int: ...
    def __getitem__(self, index: int) -> T: ...
    def __setitem__(self, index: int, val: T) -> None: ...
    def __delitem__(self, index: int) -> None: ...
    def __iter__(self) -> Iterator[T]: ...
    def __contains__(self, item: object) -> bool: ...
    def conj_mut(self, val: T) -> TransientVector[T]: ...
    def assoc_mut(self, index: int, val: T) -> TransientVector[T]: ...
    def pop_mut(self) -> TransientVector[T]: ...
    def persistent(self) -> Vector[T]: ...
    def append(self, val: T) -> TransientVector[T]: ...
    def extend(self, iterable: Any) -> TransientVector[T]: ...

# =============================================================================
# DoubleVector - Specialized vector for float64
# =============================================================================

class DoubleVector:
    def __init__(self, *args: float) -> None: ...
    def __iter__(self) -> Iterator[float]: ...
    def __len__(self) -> int: ...
    def __getitem__(self, index: int) -> float: ...
    def __hash__(self) -> int: ...
    def __eq__(self, other: object) -> bool: ...
    def __ne__(self, other: object) -> bool: ...
    def __repr__(self) -> str: ...
    def __class_getitem__(cls, item: Any) -> Any: ...
    def nth(self, index: int, default: float = ...) -> float: ...
    def conj(self, val: float) -> DoubleVector: ...
    def transient(self) -> TransientDoubleVector: ...

class TransientDoubleVector:
    def __len__(self) -> int: ...
    def __getitem__(self, index: int) -> float: ...
    def conj_mut(self, val: float) -> TransientDoubleVector: ...
    def persistent(self) -> DoubleVector: ...

# =============================================================================
# IntVector - Specialized vector for int64
# =============================================================================

class IntVector:
    def __init__(self, *args: int) -> None: ...
    def __iter__(self) -> Iterator[int]: ...
    def __len__(self) -> int: ...
    def __getitem__(self, index: int) -> int: ...
    def __hash__(self) -> int: ...
    def __eq__(self, other: object) -> bool: ...
    def __ne__(self, other: object) -> bool: ...
    def __repr__(self) -> str: ...
    def __class_getitem__(cls, item: Any) -> Any: ...
    def nth(self, index: int, default: int = ...) -> int: ...
    def conj(self, val: int) -> IntVector: ...
    def transient(self) -> TransientIntVector: ...

class TransientIntVector:
    def __len__(self) -> int: ...
    def __getitem__(self, index: int) -> int: ...
    def conj_mut(self, val: int) -> TransientIntVector: ...
    def persistent(self) -> IntVector: ...

# =============================================================================
# Map - Persistent hash map (HAMT-based)
# =============================================================================

class Map(Generic[K, V]):
    def __init__(self, *args: Any) -> None: ...
    def __iter__(self) -> Iterator[K]: ...
    def __len__(self) -> int: ...
    def __getitem__(self, key: K) -> V: ...
    def __contains__(self, key: object) -> bool: ...
    def __hash__(self) -> int: ...
    def __eq__(self, other: object) -> bool: ...
    def __ne__(self, other: object) -> bool: ...
    def __repr__(self) -> str: ...
    def __class_getitem__(cls, item: Any) -> Any: ...
    def get(self, key: K, default: V = ...) -> V: ...
    def assoc(self, key: K, val: V) -> Map[K, V]: ...
    def dissoc(self, key: K) -> Map[K, V]: ...
    def keys(self) -> Iterator[K]: ...
    def values(self) -> Iterator[V]: ...
    def items(self) -> Iterator[tuple[K, V]]: ...
    def transient(self) -> TransientMap[K, V]: ...
    def to_seq(self) -> Cons[tuple[K, V]]: ...
    def copy(self) -> Map[K, V]: ...

class TransientMap(Generic[K, V]):
    def __len__(self) -> int: ...
    def __getitem__(self, key: K) -> V: ...
    def __setitem__(self, key: K, val: V) -> None: ...
    def __delitem__(self, key: K) -> None: ...
    def __iter__(self) -> Iterator[K]: ...
    def __contains__(self, key: object) -> bool: ...
    def get(self, key: K, default: V = ...) -> V: ...
    def keys(self) -> Iterator[K]: ...
    def values(self) -> Iterator[V]: ...
    def items(self) -> Iterator[tuple[K, V]]: ...
    def assoc_mut(self, key: K, val: V) -> TransientMap[K, V]: ...
    def dissoc_mut(self, key: K) -> TransientMap[K, V]: ...
    def persistent(self) -> Map[K, V]: ...

# =============================================================================
# Set - Persistent hash set
# =============================================================================

class Set(Generic[T]):
    def __init__(self, *args: T) -> None: ...
    def __iter__(self) -> Iterator[T]: ...
    def __len__(self) -> int: ...
    def __contains__(self, item: object) -> bool: ...
    def __hash__(self) -> int: ...
    def __eq__(self, other: object) -> bool: ...
    def __ne__(self, other: object) -> bool: ...
    def __repr__(self) -> str: ...
    def __or__(self, other: Set[T]) -> Set[T]: ...
    def __and__(self, other: Set[T]) -> Set[T]: ...
    def __sub__(self, other: Set[T]) -> Set[T]: ...
    def __xor__(self, other: Set[T]) -> Set[T]: ...
    def __class_getitem__(cls, item: Any) -> Any: ...
    def conj(self, val: T) -> Set[T]: ...
    def disj(self, val: T) -> Set[T]: ...
    def transient(self) -> TransientSet[T]: ...
    def to_seq(self) -> Cons[T]: ...
    def copy(self) -> Set[T]: ...
    def isdisjoint(self, other: Any) -> bool: ...

class TransientSet(Generic[T]):
    def __len__(self) -> int: ...
    def __iter__(self) -> Iterator[T]: ...
    def __contains__(self, item: object) -> bool: ...
    def conj_mut(self, val: T) -> TransientSet[T]: ...
    def disj_mut(self, val: T) -> TransientSet[T]: ...
    def persistent(self) -> Set[T]: ...
    def add(self, val: T) -> None: ...
    def discard(self, val: T) -> None: ...
    def remove(self, val: T) -> None: ...
    def clear(self) -> None: ...

# =============================================================================
# Empty collection constants
# =============================================================================

EMPTY_VECTOR: Vector[Any]
EMPTY_DOUBLE_VECTOR: DoubleVector
EMPTY_LONG_VECTOR: IntVector
EMPTY_MAP: Map[Any, Any]
EMPTY_SET: Set[Any]

# =============================================================================
# Constructor functions
# =============================================================================

def cons(first: T, rest: Cons[T] | None = None) -> Cons[T]: ...
def vec(*args: T) -> Vector[T]: ...
def vec_f64(*args: float) -> DoubleVector: ...
def vec_i64(*args: int) -> IntVector: ...
def hash_map(*args: Any) -> Map[Any, Any]: ...
def hash_set(*args: T) -> Set[T]: ...
