from typing import Protocol, TypeVar, Any
from collections.abc import Callable
from concurrent.futures import Future

# NOTE: we have circular dependencies here,
# many of the return types from ExecutorProtocol methods themselves inherit
# ExecutorProtocol and hence import this file. Nevertheless it seems to work,
# and using string annotations sometimes makes things *not* work (mostly
# tested in vscode/pylance). Not sure why, maybe .pyi works differently from
# other imports - can't find info about it.

from .cancel_on_shutdown import CancelOnShutdownExecutor, TypedCancelOnShutdownExecutor
from .retry import RetryExecutor, TypedRetryExecutor
from .map import MapExecutor
from .flat_map import FlatMapExecutor
from .timeout import TimeoutExecutor, TypedTimeoutExecutor
from .throttle import TypedThrottleExecutor, ThrottleExecutor
from .poll import PollDescriptor, PollExecutor
from .asyncio import AsyncioExecutor

A = TypeVar("A")
B = TypeVar("B")
C = TypeVar("C")
T = TypeVar("T")

class ExecutorProtocol(Protocol):
    # Protocol for typical executor where submit output depends on the
    # return type of callables.
    def submit(self, fn: Callable[..., T], *args, **kwargs) -> Future[T]: ...
    def bind(self, *args, **kwargs): ...
    def flat_bind(self, *args, **kwargs): ...
    def with_retry(self, *args, **kwargs) -> RetryExecutor: ...
    def with_map(self, fn: Callable[[A], B], *args, **kwargs) -> MapExecutor[A, B]: ...
    def with_flat_map(
        self, fn: Callable[[A], Future[B]], *args, **kwargs
    ) -> FlatMapExecutor[A, B]: ...
    def with_poll(
        self, poll_fn: Callable[[list[PollDescriptor[A, B]]], Any], *args, **kwargs
    ) -> PollExecutor[A, B]: ...
    def with_timeout(self, *args, **kwargs) -> TimeoutExecutor: ...
    def with_throttle(self, *args, **kwargs) -> ThrottleExecutor: ...
    def with_cancel_on_shutdown(self, *args, **kwargs) -> CancelOnShutdownExecutor: ...
    def with_asyncio(self, *args, **kwargs) -> AsyncioExecutor: ...
    def __exit__(self, *args, **kwargs) -> None: ...
    def __enter__(self) -> ExecutorProtocol: ...

class TypedExecutorProtocol(Protocol[A, B]):
    # Protocol for executor where callable and output types are known.
    def submit(self, fn: Callable[..., A], *args, **kwargs) -> Future[B]: ...
    def with_retry(self, *args, **kwargs) -> TypedRetryExecutor[A, B]: ...
    def with_map(self, fn: Callable[[B], C], *args, **kwargs) -> MapExecutor[A, C]: ...
    def with_flat_map(
        self, fn: Callable[[B], Future[C]], *args, **kwargs
    ) -> FlatMapExecutor[A, C]: ...
    def with_poll(
        self, poll_fn: Callable[[list[PollDescriptor[B, C]]], Any], *args, **kwargs
    ) -> PollExecutor[A, C]: ...
    def with_timeout(self, *args, **kwargs) -> TypedTimeoutExecutor[A, B]: ...
    def with_throttle(self, *args, **kwargs) -> TypedThrottleExecutor[A, B]: ...
    def with_cancel_on_shutdown(
        self, *args, **kwargs
    ) -> TypedCancelOnShutdownExecutor[A, B]: ...
    def with_asyncio(self, *args, **kwargs) -> AsyncioExecutor: ...
    def __exit__(self, *args, **kwargs) -> None: ...
    def __enter__(self) -> TypedExecutorProtocol[A, B]: ...
