"""Helper functions for retrying."""

import asyncio
import functools
from typing import Any, Callable

import backoff
from stream.exceptions import RateLimitReached


__all__ = ["on_ratelimit_reached"]


def _get_sync_wrapper(method: Callable, decorator: Callable) -> Callable:
    @functools.wraps(method)
    def wrapper(self, *method_args, **method_kwargs) -> Any:
        decorated = decorator(self, *method_args, **method_kwargs)(method)
        return decorated(self, *method_args, **method_kwargs)

    return wrapper


def _get_async_wrapper(method: Callable, decorator: Callable) -> Callable:
    @functools.wraps(method)
    async def wrapper(self, *method_args, **method_kwargs) -> Any:
        decorated = decorator(self, *method_args, **method_kwargs)(method)
        return await decorated(self, *method_args, **method_kwargs)

    return wrapper


def _get_decorator(method: Callable, handler: Callable) -> Callable:
    """Retry on rate limit reached."""
    if asyncio.iscoroutinefunction(method):
        wrapper_gen = _get_async_wrapper
    else:
        wrapper_gen = _get_sync_wrapper
    wrapper = wrapper_gen(method, handler)

    return functools.wraps(method)(wrapper)


def _on_ratelimit_reached(self, *args, **kwargs):
    return backoff.on_exception(
        backoff.runtime,
        RateLimitReached,
        value=lambda e: self.ratelimit_info.reset_in_seconds,
        jitter=None,
        max_tries=3,
    )


on_ratelimit_reached = functools.partial(_get_decorator, handler=_on_ratelimit_reached)