"""NL2SQL 查询配额、并发和频率限制。"""
from __future__ import annotations

import time


class QueryLimiter:
    """基于 Redis 计数器的请求限制器，Redis 故障时允许降级放行。"""

    def __init__(self, redis, *, rate_window_seconds: int = 60):
        self.redis = redis
        self.rate_window_seconds = rate_window_seconds

    async def acquire(
        self,
        user_id: int,
        *,
        daily_quota: int = 0,
        max_concurrent: int = 1,
        rate_limit: int = 0,
    ) -> bool:
        """检查并占用一次配额、并发和频率额度，0 表示不限制。"""
        try:
            if daily_quota:
                daily_key = f"nl2sql:daily:{user_id}:{time.strftime('%Y%m%d')}"
                daily_count = await self.redis.incr(daily_key)
                await self.redis.expire(daily_key, 86400)
                if daily_count > daily_quota:
                    await self.redis.decr(daily_key)
                    return False
            concurrent_key = f"nl2sql:concurrent:{user_id}"
            concurrent_count = await self.redis.incr(concurrent_key)
            await self.redis.expire(concurrent_key, 3600)
            if max_concurrent and concurrent_count > max_concurrent:
                await self.redis.decr(concurrent_key)
                if daily_quota:
                    await self.redis.decr(daily_key)
                return False
            if rate_limit:
                rate_key = f"nl2sql:rate:{user_id}:{int(time.time()) // self.rate_window_seconds}"
                rate_count = await self.redis.incr(rate_key)
                await self.redis.expire(rate_key, self.rate_window_seconds + 1)
                if rate_count > rate_limit:
                    await self.redis.decr(rate_key)
                    await self.redis.decr(concurrent_key)
                    if daily_quota:
                        await self.redis.decr(daily_key)
                    return False
            return True
        except Exception:  # noqa: BLE001 Redis 故障时降级为不限制
            return True

    async def release(self, user_id: int) -> None:
        """释放一次并发额度，失败时静默降级。"""
        try:
            await self.redis.decr(f"nl2sql:concurrent:{user_id}")
        except Exception:  # noqa: BLE001 Redis 故障不能影响查询结果
            return