"""多渠道网关 Redis 滑动窗口限流。"""

from __future__ import annotations

import time
import uuid
from typing import TYPE_CHECKING

from redis.asyncio import Redis

from yuxi.channel.constants import InboundRejectionReason, channel_rate_limit_key

from .policy import SecurityCheckResult

if TYPE_CHECKING:
    from .registry import SecurityContext


class RateLimiter:
    """基于 Redis 滑动窗口的限流器。

    使用 Lua 脚本保证 "检查窗口 -> 计数 -> 写入" 的原子性，
    避免超限请求仍被计入窗口。
    """

    name = "rate_limit"
    default_priority = 400

    _ALLOW_SCRIPT = """
    local key = KEYS[1]
    local window_start = tonumber(ARGV[1])
    local max_requests = tonumber(ARGV[2])
    local member_score = tonumber(ARGV[3])
    local member = ARGV[4]
    local window_seconds = tonumber(ARGV[5])

    redis.call('zremrangebyscore', key, 0, window_start)
    local current = redis.call('zcard', key)
    if current >= max_requests then
        return 0
    end

    redis.call('zadd', key, member_score, member)
    redis.call('expire', key, window_seconds)
    return 1
    """

    def __init__(self, redis: Redis):
        self.redis = redis

    async def check(self, ctx: SecurityContext) -> SecurityCheckResult | None:
        """SecurityChecker 协议入口，保持原 is_allowed 方法可用。"""
        max_requests = self._resolve_max_requests(ctx)
        actor_id = ctx.resolved_sender_id or ctx.inbound.sender_id or ctx.inbound.peer_id or ""
        rate_key = channel_rate_limit_key(
            ctx.inbound.channel_type,
            ctx.inbound.account_id or "",
            actor_id,
        )
        if await self.is_allowed(rate_key, max_requests, window_seconds=60):
            return None
        return SecurityCheckResult(
            allowed=False,
            reason=InboundRejectionReason.RATE_LIMITED,
        )

    def _resolve_max_requests(self, ctx: SecurityContext) -> int:
        resolve = getattr(ctx.plugin, "resolve_rate_limit_policy", None)
        if resolve is not None:
            rate_policy = resolve(ctx.config, ctx.inbound.account_id)
            if rate_policy is not None:
                return rate_policy.max_requests_per_minute
        return int(ctx.config.get("rate_limit", {}).get("max_requests_per_minute", 60))

    async def is_allowed(
        self,
        key: str,
        max_requests: int,
        window_seconds: int,
    ) -> bool:
        """原有公开方法，继续可用。"""
        now = time.time()
        window_start = now - window_seconds
        member = f"{now}:{uuid.uuid4().hex}"

        allowed = await self.redis.eval(
            self._ALLOW_SCRIPT,
            1,
            key,
            window_start,
            max_requests,
            now,
            member,
            window_seconds,
        )
        return bool(allowed)