from __future__ import annotations

import random
import time
import unittest
from collections.abc import Callable
from typing import ParamSpec, TypeVar

P = ParamSpec("P")
R = TypeVar("R")


def retry_with_backoff(
    func: Callable[P, R],
    *args: P.args,
    exceptions: tuple[type[BaseException], ...] = (Exception,),
    max_attempts: int = 3,
    initial_delay: float = 0.25,
    multiplier: float = 2.0,
    max_delay: float = 10.0,
    jitter: float = 0.0,
    sleep: Callable[[float], None] = time.sleep,
    **kwargs: P.kwargs,
) -> R:
    """Run func with retries using exponential backoff."""
    if max_attempts < 1:
        raise ValueError("max_attempts must be at least 1")
    if initial_delay < 0 or max_delay < 0:
        raise ValueError("delays must be non-negative")
    if multiplier < 1:
        raise ValueError("multiplier must be at least 1")
    if not 0 <= jitter <= 1:
        raise ValueError("jitter must be between 0 and 1")

    delay = initial_delay

    for attempt in range(1, max_attempts + 1):
        try:
            return func(*args, **kwargs)
        except exceptions:
            if attempt == max_attempts:
                raise

            wait = min(delay, max_delay)

            # Jitter spreads retries out when many callers fail at once.
            if jitter:
                spread = wait * jitter
                wait = random.uniform(max(0.0, wait - spread), wait + spread)

            sleep(wait)
            delay *= multiplier

    raise RuntimeError("unreachable retry state")


class RetryWithBackoffTests(unittest.TestCase):
    def test_retries_until_success(self) -> None:
        attempts = 0
        sleeps: list[float] = []

        def flaky() -> str:
            nonlocal attempts
            attempts += 1
            if attempts < 3:
                raise TimeoutError("temporary failure")
            return "ok"

        result = retry_with_backoff(
            flaky,
            exceptions=(TimeoutError,),
            max_attempts=4,
            initial_delay=0.5,
            multiplier=2.0,
            sleep=sleeps.append,
        )

        self.assertEqual(result, "ok")
        self.assertEqual(attempts, 3)
        self.assertEqual(sleeps, [0.5, 1.0])

    def test_raises_after_final_attempt(self) -> None:
        sleeps: list[float] = []

        def always_fails() -> None:
            raise ConnectionError("still down")

        with self.assertRaises(ConnectionError):
            retry_with_backoff(
                always_fails,
                exceptions=(ConnectionError,),
                max_attempts=3,
                initial_delay=1.0,
                max_delay=1.5,
                sleep=sleeps.append,
            )

        self.assertEqual(sleeps, [1.0, 1.5])


if __name__ == "__main__":
    unittest.main()