from __future__ import annotations

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

P = ParamSpec("P")
T = TypeVar("T")


def retry_with_backoff(
    operation: Callable[P, T],
    *args: P.args,
    attempts: int = 3,
    initial_delay: float = 0.1,
    multiplier: float = 2.0,
    max_delay: float | None = None,
    retry_on: tuple[type[BaseException], ...] = (Exception,),
    sleep: Callable[[float], None] = time.sleep,
    **kwargs: P.kwargs,
) -> T:
    """Call an operation, retrying failures with exponential backoff."""
    if attempts < 1:
        raise ValueError("attempts must be at least 1")
    if initial_delay < 0:
        raise ValueError("initial_delay cannot be negative")
    if multiplier < 1:
        raise ValueError("multiplier must be at least 1")
    if max_delay is not None and max_delay < 0:
        raise ValueError("max_delay cannot be negative")

    delay = initial_delay
    for attempt in range(attempts):
        try:
            return operation(*args, **kwargs)
        except retry_on:
            if attempt == attempts - 1:
                raise

            sleep(delay)
            delay *= multiplier
            if max_delay is not None:
                delay = min(delay, max_delay)

    raise AssertionError("unreachable")


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

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

        result = retry_with_backoff(
            flaky,
            attempts=4,
            initial_delay=0.5,
            sleep=delays.append,
        )

        self.assertEqual(result, "ok")
        self.assertEqual(calls, 3)
        self.assertEqual(delays, [0.5, 1.0])

    def test_reraises_after_final_attempt(self) -> None:
        def always_fails() -> None:
            raise ValueError("permanent failure")

        with self.assertRaisesRegex(ValueError, "permanent failure"):
            retry_with_backoff(
                always_fails,
                attempts=2,
                initial_delay=0,
                sleep=lambda _: None,
            )


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