import asyncio
import inspect
from apollo_utils.service.exceptions import APIError
from http import HTTPStatus
from typing import Any, Callable, Coroutine

from server import config


def retry(
    count: int = config.DEFAULT_RETRY_COUNT,
    wait_rate: int = config.DEFAULT_RETRY_WAIT,
    max_timeout: int = config.DEFAULT_RETRY_MAX_TIMEOUT,
    auth_handler: Callable[[Any, APIError], Coroutine] = None,
):
    """Retry decorator.

    :param count: Retry count.
    :param wait_rate: Wait multiplier (current attempt number * wait_rate seconds).
    :param max_timeout: Max wait timeout.
    :param auth_handler: Auth error handler.
    """

    def inner(f: Callable):
        async def wrapped(self, *args, **kwargs):
            wait_time = 0
            for i in range(count + 1):
                try:
                    return await f(self, *args, **kwargs)
                except APIError as e:
                    # exit if it is the last try
                    if i == count:
                        raise

                    status_code = e.original_status_code or e.status_code

                    # handle rate limits
                    if status_code == HTTPStatus.TOO_MANY_REQUESTS:
                        retry_after = int(e.headers.get("Retry-After", 1))
                        wait_time = wait_time + retry_after
                        # exit if it is needed to wait too much
                        if wait_time > max_timeout:
                            raise
                        await asyncio.sleep(retry_after)
                    # handle auth errors
                    elif status_code in (HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN, 419):
                        # try to handle this only first time
                        if i > 0 or not auth_handler:
                            raise
                        await auth_handler(self, e)
                    # exit if it is an error where retry doesn't help
                    elif status_code < 500:
                        raise
                    else:
                        # otherwise wait and retry
                        current_time = wait_rate * (i + 1)
                        wait_time = wait_time + current_time
                        if wait_time > max_timeout:
                            raise
                        await asyncio.sleep(current_time)

        wrapped.__signature__ = inspect.signature(f)
        return wrapped

    return inner