office_translator/middleware/tier_quota.py

"""
Tier-based monthly translation quota.
Uses Redis counter per user per month; fallback in-memory when Redis unavailable.
Coexists with IP-based rate limiting in rate_limiting.py.

Free tier: 2 translations per calendar month.
Paid tiers (starter/pro/business/enterprise): unlimited.
"""
from __future__ import annotations

import os
import logging
from dataclasses import dataclass
from datetime import datetime, timezone, timedelta
from typing import Optional

logger = logging.getLogger(__name__)

# Free tier: 2 translations per calendar month.
FREE_TIER_MONTHLY_LIMIT = 2
KEY_PREFIX = "quota:monthly"


def _utc_month_str(dt: Optional[datetime] = None) -> str:
    """Current month in UTC as YYYY-MM."""
    t = dt or datetime.now(timezone.utc)
    return t.strftime("%Y-%m")


def _next_month_utc(dt: Optional[datetime] = None) -> datetime:
    """First second of next month UTC."""
    now = dt or datetime.now(timezone.utc)
    if now.month == 12:
        return datetime(now.year + 1, 1, 1, tzinfo=timezone.utc)
    return datetime(now.year, now.month + 1, 1, tzinfo=timezone.utc)


def _seconds_until_next_month(dt: Optional[datetime] = None) -> int:
    """Seconds until start of next month UTC."""
    now = dt or datetime.now(timezone.utc)
    return max(0, int((_next_month_utc(now) - now).total_seconds()))


@dataclass
class QuotaResult:
    """Result of a quota check."""
    allowed: bool
    remaining: int  # -1 for paid tiers (unlimited)
    reset_at_utc: datetime
    current_usage: int = 0
    limit: int = FREE_TIER_MONTHLY_LIMIT


# ---------------------------------------------------------------------------
# Redis backend
# ---------------------------------------------------------------------------


def _get_async_redis():
    """Return async Redis client or None. Uses shared client from core.redis."""
    from core.redis import get_async_redis
    return get_async_redis()


# ---------------------------------------------------------------------------
# In-memory fallback
# ---------------------------------------------------------------------------

_memory_usage: dict[tuple[str, str], int] = {}  # (user_id, month_str) -> count


def _memory_get(user_id: str, month_str: str) -> int:
    return _memory_usage.get((user_id, month_str), 0)


def _memory_incr(user_id: str, month_str: str) -> int:
    key = (user_id, month_str)
    _memory_usage[key] = _memory_usage.get(key, 0) + 1
    return _memory_usage[key]


# ---------------------------------------------------------------------------
# TierQuotaService
# ---------------------------------------------------------------------------

class TierQuotaService:
    """
    Monthly translation quota per user by tier.
    Redis key pattern: quota:monthly:{user_id}:{YYYY-MM}, TTL 32 days.
    """

    def __init__(self):
        self._redis = None

    def _redis_client(self):
        if self._redis is None:
            self._redis = _get_async_redis()
        return self._redis

    async def check_quota(self, user_id: str, tier: str) -> QuotaResult:
        """Check monthly quota. Free = 2/month; paid = unlimited."""
        reset_at = _next_month_utc()
        tier_lower = (tier or "free").lower()
        if tier_lower in ("pro", "business", "enterprise", "starter"):
            return QuotaResult(
                allowed=True,
                remaining=-1,
                reset_at_utc=reset_at,
                current_usage=0,
                limit=0,
            )
        # Free tier — monthly counter
        month_str = _utc_month_str()
        redis_client = self._redis_client()
        if redis_client:
            try:
                key = f"{KEY_PREFIX}:{user_id}:{month_str}"
                count = await redis_client.get(key)
                count = int(count or 0)
            except Exception as e:
                logger.warning("Tier quota Redis get failed: %s, using in-memory", e)
                count = _memory_get(user_id, month_str)
        else:
            count = _memory_get(user_id, month_str)
        remaining = max(0, FREE_TIER_MONTHLY_LIMIT - count)
        return QuotaResult(
            allowed=count < FREE_TIER_MONTHLY_LIMIT,
            remaining=remaining,
            reset_at_utc=reset_at,
            current_usage=count,
            limit=FREE_TIER_MONTHLY_LIMIT,
        )

    async def increment_on_success(self, user_id: str) -> None:
        """Increment monthly translation count (call after successful translation)."""
        month_str = _utc_month_str()
        redis_client = self._redis_client()
        if redis_client:
            try:
                key = f"{KEY_PREFIX}:{user_id}:{month_str}"
                pipe = redis_client.pipeline()
                pipe.incr(key)
                pipe.expire(key, 32 * 24 * 3600)  # 32 days
                await pipe.execute()
                return
            except Exception as e:
                logger.warning("Tier quota Redis increment failed: %s, using in-memory", e)
        _memory_incr(user_id, month_str)

    def seconds_until_reset(self) -> int:
        """Seconds until start of next month UTC."""
        return _seconds_until_next_month()


# Singleton
tier_quota_service = TierQuotaService()