import asyncio
import time
from typing import Awaitable, Callable, cast

from grpc import Call, RpcError, StatusCode  # type: ignore
from grpc.aio import AioRpcError  # type: ignore
from typing_extensions import ParamSpec, TypeVar

from weaviate.exceptions import WeaviateRetryError
from weaviate.logger import logger

P = ParamSpec("P")
T = TypeVar("T")


class _Retry:
    def __init__(self, n: float = 4) -> None:
        self.n = n

    async def awith_exponential_backoff(
        self,
        count: int,
        error: str,
        f: Callable[P, Awaitable[T]],
        *args: P.args,
        **kwargs: P.kwargs,
    ) -> T:
        try:
            return await f(*args, **kwargs)
        except AioRpcError as e:
            if e.code() != StatusCode.UNAVAILABLE:
                raise e
            logger.info(
                f"{error} received exception: {e}. Retrying with exponential backoff in {2**count} seconds"
            )
            await asyncio.sleep(2**count)
            if count > self.n:
                raise WeaviateRetryError(str(e), count) from e
            return await self.awith_exponential_backoff(count + 1, error, f, *args, **kwargs)

    def with_exponential_backoff(
        self,
        count: int,
        error: str,
        f: Callable[P, T],
        *args: P.args,
        **kwargs: P.kwargs,
    ) -> T:
        try:
            return f(*args, **kwargs)
        except RpcError as e:
            err = cast(Call, e)
            if err.code() != StatusCode.UNAVAILABLE:
                raise e
            logger.info(
                f"{error} received exception: {e}. Retrying with exponential backoff in {2**count} seconds"
            )
            time.sleep(2**count)
            if count > self.n:
                raise WeaviateRetryError(str(e), count) from e
            return self.with_exponential_backoff(count + 1, error, f, *args, **kwargs)
