'''Bounded-retry decorator for transient failures (network resets, partial reads, etc).
'''

import functools
import sys
import threading
import time


_BACKOFF_KINDS = ('fixed', 'linear', 'exponential')

# Thread-local current attempt number; consumed by callees (e.g. gopen_s3) to
# enable extra diagnostics on retries. Defaults to 1 outside any retry context.
retry_state = threading.local()


def _compute_sleep(base_sleep, attempt, backoff):
    if backoff == 'fixed':
        return base_sleep
    if backoff == 'linear':
        return base_sleep * attempt
    if backoff == 'exponential':
        return base_sleep * (2 ** (attempt - 1))
    raise ValueError(f'backoff must be one of {_BACKOFF_KINDS}, got {backoff!r}')


def retry(attempts=3, base_sleep=2.0, backoff='linear', exc=(Exception,), log_warning=True):
    '''Decorator: retry the wrapped call up to N times on listed exceptions.

    Sleeps between attempts (configurable backoff). Re-raises the LAST exception
    after all attempts exhaust. Non-listed exception types propagate immediately
    without retry. Per-call only; no cross-call state.

    :param attempts: int >= 1. Total attempts (including first).
    :param base_sleep: float seconds. Per-attempt sleep depends on backoff kind.
    :param backoff: 'fixed' | 'linear' | 'exponential'. Linear default mirrors
        custom/utils/aws.py + anydata/sync/sync_utils.py conventions.
    :param exc: tuple of exception classes to catch. Other exceptions bubble up
        without retry. Default = (Exception,) for unrestricted retry; restrict
        explicitly for non-idempotent or fast-fail conditions.
    :param log_warning: emit a warning between attempts (stderr).
    '''
    assert attempts >= 1, attempts
    assert backoff in _BACKOFF_KINDS, backoff

    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            last_exc = None
            for attempt in range(1, attempts + 1):
                retry_state.attempt = attempt
                try:
                    return fn(*args, **kwargs)
                except exc as e:
                    last_exc = e
                    if attempt >= attempts:
                        raise
                    sleep_s = _compute_sleep(base_sleep, attempt, backoff)
                    if log_warning:
                        print(
                            f'[retry] {fn.__name__} attempt {attempt}/{attempts} '
                            f'failed ({type(e).__name__}: {e}); '
                            f'sleeping {sleep_s:.1f}s', file=sys.stderr, flush=True)
                    time.sleep(sleep_s)
            # Unreachable, but keep static checkers happy.
            raise last_exc  # pragma: no cover

        return wrapper

    return decorator