import httpx
from loguru import logger
from tenacity import retry, retry_if_result, stop_after_attempt, wait_exponential

from httpx_limiter import AsyncRateLimitedTransport, Rate
from httpx_limiter.aiolimiter import AiolimiterAsyncLimiter

from bb_integrations_lib.secrets import AnyCredential


def log_429_retry(retry_state):
    wait = retry_state.next_action.sleep
    logger.warning(f"Got 429, retrying in {wait:.1f}s (attempt {retry_state.attempt_number})")


class BaseAPI(httpx.AsyncClient):
    """
    I'm choosing to keep our BaseAPI client to add global \
    utility methods with ease in the future. \
    A good example of this is the rate limit.
    Wrap other API clients with this.
    """

    def __init__(
        self,
        raise_errors: bool = True,
        rate_limit: Rate | None = None,
        retry_on_429: bool = True,
        max_429_retries: int = 3,
    ):
        kwargs = {}
        if rate_limit is not None:
            limiter = AiolimiterAsyncLimiter.create(rate_limit)
            kwargs["transport"] = AsyncRateLimitedTransport.create(limiter=limiter)
        super().__init__(**kwargs)
        self.raise_errors = raise_errors
        self.retry_on_429 = retry_on_429
        self.max_429_retries = max_429_retries

    @classmethod
    def from_credential(cls, credential: AnyCredential) -> "BaseAPI":
        raise NotImplementedError()

    async def send(self, request: httpx.Request, **kwargs) -> httpx.Response:
        if not self.retry_on_429:
            return await super().send(request, **kwargs)

        @retry(
            retry=retry_if_result(lambda r: r.status_code == 429),
            stop=stop_after_attempt(self.max_429_retries + 1),
            wait=wait_exponential(multiplier=1, max=60),
            before_sleep=log_429_retry,
            retry_error_callback=lambda retry_state: retry_state.outcome.result(),
        )
        async def send_with_retry():
            return await super(BaseAPI, self).send(request, **kwargs)

        return await send_with_retry()

    async def close(self):
        await self.aclose()