"""Retry policy helpers for idempotent read operations."""

from __future__ import annotations

import asyncio
import random
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from threading import RLock
from typing import TypeVar

import grpc

T = TypeVar("T")


@dataclass(slots=True, frozen=True)
class RetryPolicy:
    max_attempts: int = 3
    initial_backoff_s: float = 0.05
    max_backoff_s: float = 0.5
    multiplier: float = 2.0
    jitter_ratio: float = 0.2
    max_elapsed_time_s: float | None = 2.0
    retryable_codes: frozenset[grpc.StatusCode] = field(
        default_factory=lambda: frozenset(
            {
                grpc.StatusCode.UNAVAILABLE,
                grpc.StatusCode.DEADLINE_EXCEEDED,
                grpc.StatusCode.RESOURCE_EXHAUSTED,
            }
        )
    )

    def next_delay(self, attempt: int) -> float:
        delay = self.initial_backoff_s * (self.multiplier ** max(0, attempt - 1))
        return min(delay, self.max_backoff_s)

    def next_delay_with_jitter(self, attempt: int) -> float:
        delay = self.next_delay(attempt)
        if delay <= 0:
            return 0.0
        if self.jitter_ratio <= 0:
            return delay
        low = max(0.0, 1.0 - self.jitter_ratio)
        high = 1.0 + self.jitter_ratio
        return delay * random.uniform(low, high)


@dataclass(slots=True, frozen=True)
class CircuitBreakerPolicy:
    failure_threshold: int = 8
    recovery_timeout_s: float = 30.0
    half_open_success_threshold: int = 2
    tracked_codes: frozenset[grpc.StatusCode] = field(
        default_factory=lambda: frozenset(
            {
                grpc.StatusCode.UNAVAILABLE,
                grpc.StatusCode.DEADLINE_EXCEEDED,
                grpc.StatusCode.RESOURCE_EXHAUSTED,
            }
        )
    )

    def __post_init__(self) -> None:
        if self.failure_threshold <= 0:
            raise ValueError("failure_threshold must be > 0")
        if self.recovery_timeout_s <= 0:
            raise ValueError("recovery_timeout_s must be > 0")
        if self.half_open_success_threshold <= 0:
            raise ValueError("half_open_success_threshold must be > 0")
        if not self.tracked_codes:
            raise ValueError("tracked_codes must be non-empty")


class CircuitBreaker:
    def __init__(self, policy: CircuitBreakerPolicy) -> None:
        self._policy = policy
        self._lock = RLock()
        self._state = "closed"  # closed | open | half_open
        self._failure_count = 0
        self._opened_at = 0.0
        self._half_open_success_count = 0

    def allow_call(self) -> tuple[bool, float | None]:
        with self._lock:
            now = time.monotonic()
            if self._state == "closed":
                return True, None
            if self._state == "open":
                elapsed = now - self._opened_at
                if elapsed >= self._policy.recovery_timeout_s:
                    self._state = "half_open"
                    self._half_open_success_count = 0
                    return True, None
                return False, max(self._policy.recovery_timeout_s - elapsed, 0.0)
            return True, None

    def record_success(self) -> None:
        with self._lock:
            if self._state == "half_open":
                self._half_open_success_count += 1
                if self._half_open_success_count >= self._policy.half_open_success_threshold:
                    self._state = "closed"
                    self._failure_count = 0
                    self._opened_at = 0.0
                    self._half_open_success_count = 0
                return
            if self._state == "closed":
                self._failure_count = 0

    def record_failure(self, code: grpc.StatusCode) -> None:
        if code not in self._policy.tracked_codes:
            return
        with self._lock:
            now = time.monotonic()
            if self._state == "half_open":
                self._state = "open"
                self._opened_at = now
                self._half_open_success_count = 0
                self._failure_count = self._policy.failure_threshold
                return
            if self._state == "closed":
                self._failure_count += 1
                if self._failure_count >= self._policy.failure_threshold:
                    self._state = "open"
                    self._opened_at = now

    @property
    def state(self) -> str:
        with self._lock:
            return self._state


def run_with_retry(fn: Callable[[], T], policy: RetryPolicy) -> T:
    attempt = 1
    start = time.monotonic()
    while True:
        try:
            return fn()
        except grpc.RpcError as exc:
            code = exc.code() if hasattr(exc, "code") else grpc.StatusCode.UNKNOWN
            if attempt >= policy.max_attempts or code not in policy.retryable_codes:
                raise
            delay = policy.next_delay_with_jitter(attempt)
            if policy.max_elapsed_time_s is not None:
                elapsed = time.monotonic() - start
                if elapsed + delay > policy.max_elapsed_time_s:
                    raise
            time.sleep(delay)
            attempt += 1


async def run_with_retry_async(fn: Callable[[], Awaitable[T]], policy: RetryPolicy) -> T:
    attempt = 1
    start = time.monotonic()
    while True:
        try:
            return await fn()
        except grpc.RpcError as exc:
            code = exc.code() if hasattr(exc, "code") else grpc.StatusCode.UNKNOWN
            if attempt >= policy.max_attempts or code not in policy.retryable_codes:
                raise
            delay = policy.next_delay_with_jitter(attempt)
            if policy.max_elapsed_time_s is not None:
                elapsed = time.monotonic() - start
                if elapsed + delay > policy.max_elapsed_time_s:
                    raise
            await asyncio.sleep(delay)
            attempt += 1
