anti-plagiarism/services/api/app/core/rate_limiter.py

"""Redis-based rate limiter для проверки лимитов по тарифному плану."""

import json
from datetime import datetime, timezone

import redis

from app.config import settings

# Лимиты по тарифным планам
PLAN_LIMITS: dict[str, dict[str, int | None]] = {
    "free": {
        "search_per_day": 10,
        "summarize_per_month": 3,
        "plagiarism_per_month": 1,
        "concurrent": 1,
    },
    "student": {
        "search_per_day": None,  # None = безлимит
        "summarize_per_month": 30,
        "plagiarism_per_month": 10,
        "concurrent": 2,
    },
    "premium": {
        "search_per_day": None,
        "summarize_per_month": None,
        "plagiarism_per_month": 50,
        "concurrent": 5,
    },
    "science": {
        "search_per_day": None,
        "summarize_per_month": None,
        "plagiarism_per_month": None,
        "concurrent": 10,
    },
}

# Маппинг действий на ключи лимитов
ACTION_TO_LIMIT: dict[str, tuple[str, str]] = {
    "search": ("search_per_day", "day"),
    "summarize": ("summarize_per_month", "month"),
    "plagiarism": ("plagiarism_per_month", "month"),
    "gost": ("gost_per_month", "month"),  # ГОСТ всегда разрешён
}


def get_redis_client() -> redis.Redis:
    """Создать синхронный Redis клиент."""
    return redis.from_url(settings.REDIS_URL, decode_responses=True)


def _get_period_key(period: str) -> str:
    """Получить строку периода для Redis ключа."""
    now = datetime.now(timezone.utc)
    if period == "day":
        return now.strftime("%Y-%m-%d")
    elif period == "month":
        return now.strftime("%Y-%m")
    return now.strftime("%Y-%m-%d")


def check_and_increment_limit(user_id: int, action: str, plan: str) -> dict:
    """
    Проверить лимит и инкрементировать счётчик.

    Args:
        user_id: ID пользователя
        action: Действие (search, plagiarism, summarize, gost)
        plan: Тарифный план пользователя

    Returns:
        dict с полями:
            - allowed: bool — разрешено ли действие
            - current: int — текущее количество использований
            - limit: int | None — лимит (None = безлимит)
            - remaining: int | None — осталось использований
            - reset_at: str — когда сбрасывается счётчик
    """
    limits = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"])

    if action not in ACTION_TO_LIMIT:
        # Неизвестное действие — разрешаем
        return {"allowed": True, "current": 0, "limit": None, "remaining": None}

    limit_key, period = ACTION_TO_LIMIT[action]
    limit_value = limits.get(limit_key)

    # Безлимитный план
    if limit_value is None:
        return {
            "allowed": True,
            "current": 0,
            "limit": None,
            "remaining": None,
            "reset_at": None,
        }

    period_str = _get_period_key(period)
    redis_key = f"rate:{user_id}:{action}:{period_str}"

    r = get_redis_client()

    # Атомарно инкрементировать
    pipe = r.pipeline()
    pipe.incr(redis_key)
    # Устанавливаем TTL: для дня — 86400 сек, для месяца — 32 дня
    ttl = 86400 if period == "day" else 86400 * 32
    pipe.expire(redis_key, ttl)
    results = pipe.execute()

    current = results[0]

    if current > limit_value:
        # Декрементировать обратно (не считать запрещённые)
        r.decr(redis_key)
        current -= 1
        return {
            "allowed": False,
            "current": current,
            "limit": limit_value,
            "remaining": 0,
            "reset_at": period_str,
        }

    return {
        "allowed": True,
        "current": current,
        "limit": limit_value,
        "remaining": limit_value - current,
        "reset_at": period_str,
    }


def check_concurrent_limit(user_id: int, plan: str) -> bool:
    """
    Проверить лимит одновременных задач.

    Returns:
        True если можно создать новую задачу, False если превышен лимит.
    """
    limits = PLAN_LIMITS.get(plan, PLAN_LIMITS["free"])
    max_concurrent = limits.get("concurrent", 1)

    r = get_redis_client()
    key = f"concurrent:{user_id}"
    current = r.get(key)

    return (current is None) or (int(current) < max_concurrent)


def increment_concurrent(user_id: int) -> None:
    """Увеличить счётчик одновременных задач (при создании задачи)."""
    r = get_redis_client()
    key = f"concurrent:{user_id}"
    pipe = r.pipeline()
    pipe.incr(key)
    pipe.expire(key, 3600)  # Автосброс через 1 час
    pipe.execute()


def decrement_concurrent(user_id: int) -> None:
    """Уменьшить счётчик одновременных задач (при завершении задачи)."""
    r = get_redis_client()
    key = f"concurrent:{user_id}"
    current = r.get(key)
    if current and int(current) > 0:
        r.decr(key)