import redis
import time

class FixedWindowRateLimiter:
    def __init__(self, redis_client, key_prefix, limit, window=60):
        self.redis = redis_client
        self.key_prefix = key_prefix
        self.limit = limit
        self.window = window  # 秒

    def is_allowed(self, key: str) -> bool:
        full_key = f"{self.key_prefix}:{key}"
        current = int(time.time())
        window_start = (current // self.window) * self.window
        
        # 使用 Redis pipeline 保證原子性
        pipe = self.redis.pipeline()
        pipe.zremrangebyscore(full_key, 0, window_start - 1)
        pipe.zcard(full_key)
        pipe.zadd(full_key, {str(current): current})
        pipe.expire(full_key, self.window + 1)
        _, count, _, _ = pipe.execute()
        
        return count < self.limit

class SlidingWindowRateLimiter:
    def __init__(self, redis_client, key_prefix, limit, window=60):
        self.redis = redis_client
        self.key_prefix = key_prefix
        self.limit = limit
        self.window = window

    def is_allowed(self, key: str) -> bool:
        full_key = f"{self.key_prefix}:{key}"
        now = time.time()
        window_start = now - self.window
        
        pipe = self.redis.pipeline()
        # 移除窗口外的請(qǐng)求
        pipe.zremrangebyscore(full_key, 0, window_start)
        # 獲取當(dāng)前窗口內(nèi)請(qǐng)求數(shù)
        pipe.zcard(full_key)
        # 添加當(dāng)前請(qǐng)求
        pipe.zadd(full_key, {str(now): now})
        # 設(shè)置過(guò)期時(shí)間（避免冷 key 占用內(nèi)存）
        pipe.expire(full_key, int(self.window) + 1)
        
        _, count, _, _ = pipe.execute()
        return count <= self.limit

import json

class TokenBucketRateLimiter:
    LUA_SCRIPT = """
    local tokens_key = KEYS[1]
    local timestamp_key = KEYS[2]
    local rate = tonumber(ARGV[1])
    local capacity = tonumber(ARGV[2])
    local now = tonumber(ARGV[3])
    local requested = tonumber(ARGV[4])

    local last_tokens = redis.call('GET', tokens_key)
    if not last_tokens then
        last_tokens = capacity
    end

    local last_time = redis.call('GET', timestamp_key)
    if not last_time then
        last_time = now
    end

    local tokens = tonumber(last_tokens)
    local last_time = tonumber(last_time)

    -- 計(jì)算新令牌數(shù)
    local new_tokens = tokens + (now - last_time) * rate
    if new_tokens > capacity then
        new_tokens = capacity
    end

    local allowed = new_tokens >= requested
    if allowed then
        new_tokens = new_tokens - requested
    end

    -- 更新?tīng)顟B(tài)
    redis.call('SET', tokens_key, new_tokens)
    redis.call('SET', timestamp_key, now)
    redis.call('EXPIRE', tokens_key, 10)
    redis.call('EXPIRE', timestamp_key, 10)

    return {allowed and 1 or 0, new_tokens}
    """

    def __init__(self, redis_client, key_prefix, rate, capacity):
        self.redis = redis_client
        self.key_prefix = key_prefix
        self.rate = rate      # 令牌生成速率（token/秒）
        self.capacity = capacity  # 桶容量
        self.script = self.redis.register_script(self.LUA_SCRIPT)

    def is_allowed(self, key: str, tokens=1) -> bool:
        tokens_key = f"{self.key_prefix}:tokens:{key}"
        timestamp_key = f"{self.key_prefix}:time:{key}"
        now = time.time()
        
        result = self.script(
            keys=[tokens_key, timestamp_key],
            args=[self.rate, self.capacity, now, tokens]
        )
        return bool(result[0])

# rate_limiter.py
from fastapi import Request, HTTPException, status
from starlette.middleware.base import BaseHTTPMiddleware
from .token_bucket import TokenBucketRateLimiter  # 上述實(shí)現(xiàn)

redis_client = redis.Redis(host="localhost", port=6379, decode_responses=True)
limiter = TokenBucketRateLimiter(
    redis_client, 
    key_prefix="api", 
    rate=10,      # 10 token/秒
    capacity=20   # 允許突發(fā) 20 次
)

class RateLimitMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        # 獲取限流 key（可按 IP、用戶 ID、路徑組合）
        client_ip = request.client.host
        path = request.url.path
        
        # 例如：按 IP + 路徑限流
        rate_key = f"{client_ip}:{path}"
        
        if not limiter.is_allowed(rate_key):
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail="Too Many Requests",
                headers={"Retry-After": "1"}  # 建議重試時(shí)間
            )
        
        response = await call_next(request)
        return response

# main.py
from fastapi import FastAPI
from rate_limiter import RateLimitMiddleware

app = FastAPI()
app.add_middleware(RateLimitMiddleware)

@app.get("/hello")
async def hello():
    return {"message": "Hello World"}

from functools import wraps

def rate_limit(rate: float, capacity: int, key_func=None):
    def decorator(func):
        limiter = TokenBucketRateLimiter(
            redis_client, 
            key_prefix=f"func:{func.__name__}", 
            rate=rate, 
            capacity=capacity
        )
        
        @wraps(func)
        async def wrapper(*args, **kwargs):
            # 從請(qǐng)求中提取 key（需根據(jù)框架調(diào)整）
            request = kwargs.get("request") or args[0]
            key = key_func(request) if key_func else request.client.host
            
            if not limiter.is_allowed(key):
                raise HTTPException(429, "Rate limit exceeded")
            
            return await func(*args, **kwargs)
        return wrapper
    return decorator

# 使用
@app.post("/send-sms")
@rate_limit(rate=1/60, capacity=1)  # 每用戶每分鐘 1 次
async def send_sms(request: Request, phone: str):
    # 發(fā)送短信邏輯
    pass

# 降級(jí)示例
try:
    allowed = redis_limiter.is_allowed(key)
except redis.ConnectionError:
    allowed = local_limiter.is_allowed(key)  # 本地限流

# Prometheus 指標(biāo)
from prometheus_client import Counter

RATE_LIMIT_REJECTED = Counter(
    "rate_limit_rejected_total", 
    "Total number of rate limited requests",
    ["endpoint", "client"]
)

# 在限流拒絕時(shí)增加計(jì)數(shù)
if not allowed:
    RATE_LIMIT_REJECTED.labels(endpoint=path, client=ip).inc()
    raise HTTPException(429, ...)

HTTP/1.1 429 Too Many Requests
Retry-After: 5
X-RateLimit-Limit: 100
X-RateLimit-Remaining: 0
X-RateLimit-Reset: 1708761600