#  Copyright 2026 zhaoxi826
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.


import asyncio
from functools import wraps
from kilostar.utils.error import RetryableError


def retry_on_retryable_error(max_retries=3, base_delay=1):
    """指数退避重试装饰器：仅在抛出 ``RetryableError`` 时重试。

    同步/异步函数都支持。第 n 次重试前会 ``sleep(base_delay * 2**n)``。

    Args:
        max_retries: 最多尝试次数（含首次），超过后会把最后一次异常重新抛出。
        base_delay: 退避基准秒数。
    """

    def decorator(func):
        if asyncio.iscoroutinefunction(func):

            @wraps(func)
            async def async_wrapper(*args, **kwargs):
                for attempt in range(max_retries):
                    try:
                        return await func(*args, **kwargs)
                    except RetryableError:
                        if attempt == max_retries - 1:
                            raise
                        await asyncio.sleep(base_delay * (2**attempt))

            return async_wrapper
        else:

            @wraps(func)
            def sync_wrapper(*args, **kwargs):
                import time

                for attempt in range(max_retries):
                    try:
                        return func(*args, **kwargs)
                    except RetryableError:
                        if attempt == max_retries - 1:
                            raise
                        time.sleep(base_delay * (2**attempt))

            return sync_wrapper

    return decorator