#  Copyright (C) 2023
#  ABM, Moscow
#
#  UNPUBLISHED PROPRIETARY MATERIAL.
#  ALL RIGHTS RESERVED.
#
#  Authors: Mikhail Mamrov <m.mamrov@abm-jsc.ru>

import asyncio
import itertools
import logging
from functools import wraps
from time import sleep
from typing import Callable, Any, Union, Type, Optional


logger = logging.getLogger(__name__)


def retry(func: Callable[[Any], Any] = None,
          attempts_limit: Optional[int] = None,
          sleep_time_sec: float = 5.0,
          retrying_exceptions: Union[Type[Exception], tuple[Type[Exception], ...]] = Exception,
          ) -> Callable[[Any], Any]:
    def decorator(handler_) -> Callable[[Any], Any]:
        def handle_exception(er: Exception, attempt: int) -> None:
            logger.error(f"Execution '{handler_}' unsuccessful, because: {repr(er)}. Retry: {attempt}/{attempts_limit}")
            if attempts_limit is not None and attempt >= attempts_limit:
                raise

        if asyncio.iscoroutinefunction(handler_):
            @wraps(handler_)
            async def wrapper(*args, **kwargs) -> Callable[[Any], Any]:
                for attempt in itertools.count():
                    try:
                        return await handler_(*args, **kwargs)
                    except retrying_exceptions as er:
                        handle_exception(er, attempt)
                    await asyncio.sleep(sleep_time_sec)
            return wrapper
        else:
            @wraps(handler_)
            def wrapper(*args, **kwargs) -> Callable[[Any], Any]:
                for attempt in itertools.count():
                    try:
                        return handler_(*args, **kwargs)
                    except retrying_exceptions as er:
                        handle_exception(er, attempt)
                    sleep(sleep_time_sec)
            return wrapper

    return decorator if func is None else decorator(func)
