#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PerToolBox Server - 限流中间件
Copyright (C) 2024 Sea Network Technology Studio
Author: Canglan <admin@sea-studio.top>
License: AGPL v3
"""

from fastapi import Request, HTTPException
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
from ..config import settings

limiter = Limiter(key_func=get_remote_address)

def setup_rate_limit(app):
    if settings.RATE_LIMIT_ENABLED:
        app.state.limiter = limiter
        app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

def rate_limit(requests: int = None, period: int = None):
    if not settings.RATE_LIMIT_ENABLED:
        return lambda func: func

    req = requests or settings.RATE_LIMIT_REQUESTS
    per = period or settings.RATE_LIMIT_PERIOD
    return limiter.limit(f"{req}/{per} seconds")