"""Redis state backend for multi-instance deployments."""

from __future__ import annotations

from typing import TYPE_CHECKING

from llm_rotator.backends import AbstractStateBackend

if TYPE_CHECKING:
    import redis.asyncio as aioredis


class RedisBackend(AbstractStateBackend):
    """Redis-backed state backend implementing :class:`AbstractStateBackend`.

    All keys are stored under a configurable prefix (default ``llm_rotator:``).
    """

    def __init__(
        self,
        redis: aioredis.Redis,
        prefix: str = "llm_rotator:",
    ) -> None:
        self._redis = redis
        self._prefix = prefix

    @classmethod
    async def from_url(
        cls,
        url: str,
        prefix: str = "llm_rotator:",
    ) -> RedisBackend:
        """Create a :class:`RedisBackend` from a Redis connection URL."""
        import redis.asyncio as aioredis

        client = aioredis.from_url(url, decode_responses=True)
        return cls(redis=client, prefix=prefix)

    def _key(self, *parts: str) -> str:
        return self._prefix + ":".join(parts)

    # --- Circuit Breaker ---

    async def mark_key_dead(self, key_id: str) -> None:
        """Mark a key as permanently dead (no TTL)."""
        await self._redis.set(self._key("dead", key_id), "1")

    async def block_key_for_model(
        self, key_id: str, model: str, ttl_seconds: int
    ) -> None:
        """Block a key for a specific model with TTL."""
        await self._redis.set(
            self._key("block", key_id, model), "1", ex=ttl_seconds
        )

    async def is_key_available(self, key_id: str, model: str) -> bool:
        """Check if a key is available for a specific model."""
        if await self._redis.exists(self._key("dead", key_id)):
            return False
        return not await self._redis.exists(self._key("block", key_id, model))

    # --- Quota Counters ---

    async def increment_quota(
        self, scope: str, amount: int, ttl_seconds: int
    ) -> int:
        """Atomically increment quota counter, setting TTL on first write."""
        key = self._key("quota", scope)
        pipe = self._redis.pipeline()
        pipe.incrby(key, amount)
        pipe.ttl(key)
        results = await pipe.execute()
        new_value: int = results[0]
        current_ttl: int = results[1]

        # TTL == -1 means no expiry set yet (first increment or key had no TTL)
        if current_ttl == -1:
            await self._redis.expire(key, ttl_seconds)

        return new_value

    async def get_quota_usage(self, scope: str) -> int:
        """Get current quota usage for a scope."""
        val = await self._redis.get(self._key("quota", scope))
        if val is None:
            return 0
        return int(val)
