"""Per-category retry rules from the v0.1 design, section 12."""

from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Callable

from .domain import ErrorCategory


@dataclass(frozen=True)
class RetryRule:
    max_attempts: int
    base_delay_s: float = 0.0

    def backoff_s(self, attempt: int) -> float:
        """Delay before retry number `attempt`, counting from 1."""
        if self.base_delay_s <= 0 or attempt < 1:
            return 0.0
        return self.base_delay_s * (2 ** (attempt - 1))


#: Everything absent from this table fails closed with no retry.
RULES = {
    ErrorCategory.RATE_LIMIT: RetryRule(max_attempts=4, base_delay_s=2.0),
    ErrorCategory.NETWORK: RetryRule(max_attempts=3, base_delay_s=1.0),
    ErrorCategory.TIMEOUT: RetryRule(max_attempts=1),
    ErrorCategory.PROTOCOL: RetryRule(max_attempts=1),
}

NO_RETRY = RetryRule(max_attempts=0)


@dataclass(frozen=True)
class RetryPolicy:
    #: Injected so tests never actually wait.
    sleep: Callable[[float], None] = time.sleep

    def for_category(self, category: ErrorCategory) -> RetryRule:
        return RULES.get(category, NO_RETRY)

    def should_retry(self, category: ErrorCategory, attempts_so_far: int) -> bool:
        return attempts_so_far < self.for_category(category).max_attempts

    def wait(self, category: ErrorCategory, attempts_so_far: int) -> float:
        delay = self.for_category(category).backoff_s(attempts_so_far)
        if delay > 0:
            self.sleep(delay)
        return delay