import attr
import trio
from abc import ABCMeta
from typing import (
    Any,
    AsyncContextManager,
    Awaitable,
    Callable,
    ContextManager,
    Optional,
    Tuple,
    TypeVar,
    Union,
    final,
)

_F = TypeVar("_F", bound=Callable[..., Any])

async def wait_all_tasks_blocked(cushion: float = ...) -> None: ...
def trio_test(fn: _F) -> _F: ...
@final
class MockClock(trio.abc.Clock):
    rate: float
    autojump_threshold: float
    def __init__(self, rate: float = ..., autojump_threshold: float = ...): ...
    def start_clock(self) -> None: ...
    def current_time(self) -> float: ...
    def deadline_to_sleep_time(self, deadline: float) -> float: ...
    def jump(self, seconds: float) -> None: ...

def assert_checkpoints() -> ContextManager[None]: ...
def assert_no_checkpoints() -> ContextManager[None]: ...
@final
@attr.s(eq=False, hash=False)
class Sequencer(metaclass=ABCMeta):
    def __call__(self, position: int) -> AsyncContextManager[None]: ...

_StreamMaker = Callable[
    [], Awaitable[Tuple[trio.abc.SendStream, trio.abc.ReceiveStream]]
]

async def check_one_way_stream(
    stream_maker: _StreamMaker, clogged_stream_maker: Optional[_StreamMaker]
) -> None: ...
async def check_two_way_stream(
    stream_maker: _StreamMaker, clogged_stream_maker: Optional[_StreamMaker]
) -> None: ...
async def check_half_closeable_stream(
    stream_maker: _StreamMaker, clogged_stream_maker: Optional[_StreamMaker]
) -> None: ...

_MemoryStreamHook = Optional[Callable[[], Awaitable[Any]]]
_MemoryStreamSyncHook = Optional[Callable[[], Any]]

@final
class MemorySendStream(trio.abc.SendStream):
    send_all_hook: _MemoryStreamHook = ...
    wait_send_all_might_not_block_hook: _MemoryStreamHook = ...
    close_hook: _MemoryStreamSyncHook = ...
    def __init__(
        self,
        send_all_hook: _MemoryStreamHook = ...,
        wait_send_all_might_not_block_hook: _MemoryStreamHook = ...,
        close_hook: _MemoryStreamSyncHook = ...,
    ) -> None: ...
    async def send_all(self, data: Union[bytes, bytearray, memoryview]) -> None: ...
    async def wait_send_all_might_not_block(self) -> None: ...
    async def aclose(self) -> None: ...
    def close(self) -> None: ...
    async def get_data(self, max_bytes: Optional[int] = ...) -> bytes: ...
    def get_data_nowait(self, max_bytes: Optional[int] = ...) -> bytes: ...

@final
class MemoryReceiveStream(trio.abc.ReceiveStream):
    receive_some_hook: _MemoryStreamHook = ...
    close_hook: _MemoryStreamSyncHook = ...
    def __init__(
        self,
        receive_some_hook: _MemoryStreamHook = None,
        close_hook: _MemoryStreamSyncHook = None,
    ) -> None: ...
    async def receive_some(self, max_bytes: Optional[int] = ...) -> bytes: ...
    async def aclose(self) -> None: ...
    def close(self) -> None: ...
    def put_data(self, data: bytes) -> None: ...
    def put_eof(self) -> None: ...

def memory_stream_pump(
    memory_send_stream: MemorySendStream,
    memory_receive_stream: MemoryReceiveStream,
    *,
    max_bytes: Optional[int] = ...,
) -> bool: ...
def memory_stream_one_way_pair() -> Tuple[MemorySendStream, MemoryReceiveStream]: ...
def memory_stream_pair() -> Tuple[trio.StapledStream, trio.StapledStream]: ...
def lockstep_stream_one_way_pair() -> Tuple[MemorySendStream, MemoryReceiveStream]: ...
def lockstep_stream_pair() -> Tuple[trio.StapledStream, trio.StapledStream]: ...
async def open_stream_to_socket_listener(
    socket_listener: trio.SocketListener,
) -> trio.SocketStream: ...
