"""Redis 多维度限流器（令牌桶算法）.

三维度限流（02-architecture-design.md §7）：
  - user: 每用户 10 req/min
  - IP: 每 IP 30 req/min
  - school: 每学校 100 req/min

使用 Redis Lua 脚本实现原子令牌桶，保证线程安全。
全并行模式：Redis 不可用时降级放行（记录警告）。
"""

import time
from dataclasses import dataclass

import structlog
from redis.asyncio import Redis
from redis.exceptions import RedisError

from .errors import AIRateLimitedError

logger = structlog.get_logger()

# Lua 脚本：原子令牌桶
# KEYS[1] = 限流 key
# ARGV[1] = 容量（max tokens）
# ARGV[2] = 补充速率（tokens per second）
# ARGV[3] = 当前时间戳（秒）
# ARGV[4] = 请求消耗 tokens（通常 1）
# 返回: {allowed(1/0), remaining}
TOKEN_BUCKET_LUA = """
local key = KEYS[1]
local capacity = tonumber(ARGV[1])
local refill_rate = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local cost = tonumber(ARGV[4])

local bucket = redis.call('HMGET', key, 'tokens', 'last_refill')
local tokens = tonumber(bucket[1]) or capacity
local last_refill = tonumber(bucket[2]) or now

-- 补充令牌
local elapsed = math.max(0, now - last_refill)
tokens = math.min(capacity, tokens + elapsed * refill_rate)

local allowed = 0
if tokens >= cost then
    tokens = tokens - cost
    allowed = 1
end

redis.call('HMSET', key, 'tokens', tokens, 'last_refill', now)
redis.call('EXPIRE', key, 120)

return {allowed, math.floor(tokens)}
"""


@dataclass
class RateLimitResult:
    """限流检查结果."""

    allowed: bool
    dimension: str
    limit: int
    remaining: int


class RateLimiter:
    """多维度限流器.

    三维度独立检查，任一维度超限即拒绝。
    全并行模式：Redis 不可用时降级放行。
    """

    def __init__(
        self,
        redis: Redis | None = None,
        user_limit: int = 10,
        ip_limit: int = 30,
        school_limit: int = 100,
        window_seconds: int = 60,
    ) -> None:
        self._redis = redis
        self._user_limit = user_limit
        self._ip_limit = ip_limit
        self._school_limit = school_limit
        self._window_seconds = window_seconds
        self._lua_sha: str | None = None

    async def _ensure_lua(self) -> bool:
        """确保 Lua 脚本已加载."""
        if self._redis is None:
            return False
        if self._lua_sha is not None:
            return True
        try:
            self._lua_sha = await self._redis.script_load(TOKEN_BUCKET_LUA)
            return True
        except RedisError as exc:
            logger.warning("rate_limiter_lua_load_failed", error=str(exc))
            return False

    async def check(
        self,
        user_id: str = "",
        ip: str = "",
        school_id: str = "",
    ) -> list[RateLimitResult]:
        """检查三维度限流.

        Returns:
            各维度检查结果列表

        Raises:
            AIRateLimitedError: 任一维度超限
        """
        results: list[RateLimitResult] = []

        if user_id:
            results.append(
                await self._check_dimension("user", user_id, self._user_limit),
            )
        if ip:
            results.append(
                await self._check_dimension("ip", ip, self._ip_limit),
            )
        if school_id:
            results.append(
                await self._check_dimension("school", school_id, self._school_limit),
            )

        for result in results:
            if not result.allowed:
                logger.warning(
                    "rate_limit_exceeded",
                    dimension=result.dimension,
                    limit=result.limit,
                    remaining=result.remaining,
                )
                raise AIRateLimitedError(result.dimension, result.limit)

        return results

    async def _check_dimension(
        self,
        dimension: str,
        identifier: str,
        limit: int,
    ) -> RateLimitResult:
        """检查单维度限流."""
        if not await self._ensure_lua():
            # Redis 不可用，降级放行
            logger.debug("rate_limiter_degraded_redis_unavailable", dimension=dimension)
            return RateLimitResult(
                allowed=True,
                dimension=dimension,
                limit=limit,
                remaining=limit,
            )

        key = f"ratelimit:{dimension}:{identifier}"
        refill_rate = limit / self._window_seconds
        now = time.time()

        try:
            result = await self._redis.evalsha(  # type: ignore[union-attr]
                self._lua_sha,
                1,
                key,
                limit,
                refill_rate,
                now,
                1,
            )
            allowed = bool(result[0])
            remaining = int(result[1])
            return RateLimitResult(
                allowed=allowed,
                dimension=dimension,
                limit=limit,
                remaining=remaining,
            )
        except RedisError as exc:
            logger.warning(
                "rate_limiter_check_failed_degraded",
                dimension=dimension,
                error=str(exc),
            )
            return RateLimitResult(
                allowed=True,
                dimension=dimension,
                limit=limit,
                remaining=limit,
            )
