"""State backends: AbstractStateBackend (ABC) + InMemoryBackend."""

from __future__ import annotations

import asyncio
import time
from abc import ABC, abstractmethod
from collections.abc import Callable


class AbstractStateBackend(ABC):
    """Interface for storing circuit breaker blocks and quota counters."""

    # --- Circuit Breaker ---

    @abstractmethod
    async def mark_key_dead(self, key_id: str) -> None: ...

    @abstractmethod
    async def block_key_for_model(self, key_id: str, model: str, ttl_seconds: int) -> None: ...

    @abstractmethod
    async def is_key_available(self, key_id: str, model: str) -> bool: ...

    # --- Quota Counters ---

    @abstractmethod
    async def increment_quota(self, scope: str, amount: int, ttl_seconds: int) -> int:
        """Atomically increment counter, return new value. TTL = time until reset."""
        ...

    @abstractmethod
    async def get_quota_usage(self, scope: str) -> int: ...


class InMemoryBackend(AbstractStateBackend):
    """In-memory state backend with injectable clock for testability."""

    def __init__(self, clock: Callable[[], float] | None = None) -> None:
        self._clock = clock or time.monotonic
        self._dead_keys: set[str] = set()
        self._model_blocks: dict[str, float] = {}  # "key_id:model" -> expires_at
        self._quotas: dict[str, tuple[int, float]] = {}  # scope -> (value, expires_at)
        self._lock = asyncio.Lock()

    async def mark_key_dead(self, key_id: str) -> None:
        async with self._lock:
            self._dead_keys.add(key_id)

    async def block_key_for_model(self, key_id: str, model: str, ttl_seconds: int) -> None:
        async with self._lock:
            expires_at = self._clock() + ttl_seconds
            self._model_blocks[f"{key_id}:{model}"] = expires_at

    async def is_key_available(self, key_id: str, model: str) -> bool:
        async with self._lock:
            if key_id in self._dead_keys:
                return False
            block_key = f"{key_id}:{model}"
            if block_key in self._model_blocks:
                if self._clock() < self._model_blocks[block_key]:
                    return False
                del self._model_blocks[block_key]
            return True

    async def increment_quota(self, scope: str, amount: int, ttl_seconds: int) -> int:
        async with self._lock:
            now = self._clock()
            if scope in self._quotas:
                current_value, expires_at = self._quotas[scope]
                if now >= expires_at:
                    current_value = 0
            else:
                current_value = 0
            new_value = current_value + amount
            self._quotas[scope] = (new_value, now + ttl_seconds)
            return new_value

    async def get_quota_usage(self, scope: str) -> int:
        async with self._lock:
            if scope not in self._quotas:
                return 0
            value, expires_at = self._quotas[scope]
            if self._clock() >= expires_at:
                del self._quotas[scope]
                return 0
            return value
