from __future__ import annotations

import grpc
import pytest

import kyrodb.retry as retry_module
from kyrodb.retry import (
    CircuitBreaker,
    CircuitBreakerPolicy,
    RetryPolicy,
    run_with_retry,
    run_with_retry_async,
)


class _FakeRpcError(grpc.RpcError):
    def __init__(self, code: grpc.StatusCode) -> None:
        super().__init__()
        self._code = code

    def code(self) -> grpc.StatusCode:
        return self._code


def test_run_with_retry_retries_then_succeeds() -> None:
    attempts = 0
    policy = RetryPolicy(max_attempts=3, initial_backoff_s=0.0, max_backoff_s=0.0)

    def flaky_call() -> str:
        nonlocal attempts
        attempts += 1
        if attempts < 3:
            raise _FakeRpcError(grpc.StatusCode.UNAVAILABLE)
        return "ok"

    assert run_with_retry(flaky_call, policy) == "ok"
    assert attempts == 3


def test_run_with_retry_does_not_retry_non_retryable() -> None:
    policy = RetryPolicy(max_attempts=5, initial_backoff_s=0.0, max_backoff_s=0.0)

    def call() -> str:
        raise _FakeRpcError(grpc.StatusCode.INVALID_ARGUMENT)

    with pytest.raises(_FakeRpcError):
        run_with_retry(call, policy)


@pytest.mark.asyncio
async def test_run_with_retry_async_retries_then_succeeds() -> None:
    attempts = 0
    policy = RetryPolicy(max_attempts=3, initial_backoff_s=0.0, max_backoff_s=0.0)

    async def flaky_call() -> str:
        nonlocal attempts
        attempts += 1
        if attempts < 3:
            raise _FakeRpcError(grpc.StatusCode.UNAVAILABLE)
        return "ok"

    assert await run_with_retry_async(flaky_call, policy) == "ok"
    assert attempts == 3


def test_next_delay_with_jitter_within_bounds(monkeypatch: pytest.MonkeyPatch) -> None:
    policy = RetryPolicy(
        initial_backoff_s=1.0,
        max_backoff_s=1.0,
        jitter_ratio=0.2,
        max_elapsed_time_s=None,
    )
    monkeypatch.setattr(retry_module.random, "uniform", lambda low, high: low)
    low = policy.next_delay_with_jitter(1)
    monkeypatch.setattr(retry_module.random, "uniform", lambda low, high: high)
    high = policy.next_delay_with_jitter(1)
    assert low == pytest.approx(0.8)
    assert high == pytest.approx(1.2)


def test_run_with_retry_respects_elapsed_budget(monkeypatch: pytest.MonkeyPatch) -> None:
    attempts = 0
    fake_clock = {"now": 0.0}

    def fake_sleep(delay: float) -> None:
        fake_clock["now"] += delay

    def fake_monotonic() -> float:
        return fake_clock["now"]

    monkeypatch.setattr(retry_module.time, "sleep", fake_sleep)
    monkeypatch.setattr(retry_module.time, "monotonic", fake_monotonic)

    policy = RetryPolicy(
        max_attempts=10,
        initial_backoff_s=0.1,
        max_backoff_s=0.1,
        jitter_ratio=0.0,
        max_elapsed_time_s=0.15,
    )

    def always_unavailable() -> str:
        nonlocal attempts
        attempts += 1
        raise _FakeRpcError(grpc.StatusCode.UNAVAILABLE)

    with pytest.raises(_FakeRpcError):
        run_with_retry(always_unavailable, policy)
    assert attempts == 2


def test_circuit_breaker_opens_and_recovers(monkeypatch: pytest.MonkeyPatch) -> None:
    fake_clock = {"now": 0.0}
    monkeypatch.setattr(retry_module.time, "monotonic", lambda: fake_clock["now"])

    breaker = CircuitBreaker(
        CircuitBreakerPolicy(
            failure_threshold=2,
            recovery_timeout_s=5.0,
            half_open_success_threshold=2,
        )
    )

    assert breaker.allow_call() == (True, None)
    breaker.record_failure(grpc.StatusCode.UNAVAILABLE)
    assert breaker.state == "closed"

    breaker.record_failure(grpc.StatusCode.UNAVAILABLE)
    allowed, retry_after_s = breaker.allow_call()
    assert allowed is False
    assert retry_after_s == pytest.approx(5.0)
    assert breaker.state == "open"

    fake_clock["now"] = 5.0
    assert breaker.allow_call() == (True, None)
    assert breaker.state == "half_open"

    breaker.record_success()
    assert breaker.state == "half_open"
    breaker.record_success()
    assert breaker.state == "closed"


def test_circuit_breaker_failures_in_half_open_reset_to_open(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    fake_clock = {"now": 0.0}
    monkeypatch.setattr(retry_module.time, "monotonic", lambda: fake_clock["now"])

    breaker = CircuitBreaker(
        CircuitBreakerPolicy(
            failure_threshold=1,
            recovery_timeout_s=5.0,
            half_open_success_threshold=1,
        )
    )

    # Fail to open
    breaker.record_failure(grpc.StatusCode.UNAVAILABLE)
    assert breaker.state == "open"

    # Advance time to recovery
    fake_clock["now"] = 10.0
    allowed, _ = breaker.allow_call()
    assert allowed is True
    assert breaker.state == "half_open"

    # Failure in half-open should immediately reopen
    breaker.record_failure(grpc.StatusCode.UNAVAILABLE)
    assert breaker.state == "open"

    # Verify timeout is reset
    allowed, retry_after = breaker.allow_call()
    assert allowed is False
    assert retry_after == pytest.approx(5.0)


def test_circuit_breaker_ignores_untracked_codes() -> None:
    breaker = CircuitBreaker(CircuitBreakerPolicy(failure_threshold=1))
    breaker.record_failure(grpc.StatusCode.INVALID_ARGUMENT)
    assert breaker.state == "closed"


def test_circuit_breaker_policy_validates_parameters() -> None:
    with pytest.raises(ValueError, match="failure_threshold"):
        CircuitBreakerPolicy(failure_threshold=0)
    with pytest.raises(ValueError, match="recovery_timeout_s"):
        CircuitBreakerPolicy(recovery_timeout_s=0.0)
    with pytest.raises(ValueError, match="half_open_success_threshold"):
        CircuitBreakerPolicy(half_open_success_threshold=0)
    with pytest.raises(ValueError, match="tracked_codes"):
        CircuitBreakerPolicy(tracked_codes=frozenset())
