rateLimit.ts

/**
 * Section 1: Backend Core - Rate Limiting Middleware
 * 
 * Implements token bucket algorithm to prevent API abuse
 * Tracks rate limits per user with configurable windows
 */

/**
 * Token bucket for rate limiting
 */
interface TokenBucket {
  tokens: number;
  lastRefillTime: number;
}

/**
 * Rate limit configuration
 */
export const RATE_LIMIT_CONFIG = {
  requestsPerMinute: 30,
  requestsPerHour: 500,
  burstSize: 5,
};

/**
 * In-memory store for rate limit buckets
 * In production, use Redis for distributed rate limiting
 */
const rateLimitBuckets = new Map<string, TokenBucket>();

/**
 * Clean up old buckets periodically (every 5 minutes)
 */
setInterval(() => {
  const now = Date.now();
  const fiveMinutesAgo = now - 5 * 60 * 1000;

  rateLimitBuckets.forEach((bucket, key) => {
    if (bucket.lastRefillTime < fiveMinutesAgo) {
      rateLimitBuckets.delete(key);
    }
  });
}, 5 * 60 * 1000);

/**
 * Check if a request is allowed based on rate limits
 * 
 * @param userId - User identifier (or IP for anonymous users)
 * @param requestType - Type of request (e.g., "chat", "imagine")
 * @returns Object with allowed status and remaining tokens
 */
export function checkRateLimit(
  userId: string,
  requestType: string = "default"
): { allowed: boolean; remainingTokens: number; resetIn: number } {
  const key = `${userId}:${requestType}`;
  const now = Date.now();
  const refillRate = RATE_LIMIT_CONFIG.requestsPerMinute / 60; // tokens per second

  let bucket = rateLimitBuckets.get(key);

  if (!bucket) {
    bucket = {
      tokens: RATE_LIMIT_CONFIG.burstSize,
      lastRefillTime: now,
    };
    rateLimitBuckets.set(key, bucket);
  }

  // Calculate elapsed time since last refill
  const elapsedSeconds = (now - bucket.lastRefillTime) / 1000;

  // Refill tokens based on elapsed time
  const tokensToAdd = elapsedSeconds * refillRate;
  bucket.tokens = Math.min(
    RATE_LIMIT_CONFIG.burstSize,
    bucket.tokens + tokensToAdd
  );
  bucket.lastRefillTime = now;

  // Check if request is allowed
  const allowed = bucket.tokens >= 1;

  if (allowed) {
    bucket.tokens -= 1;
  }

  // Calculate reset time (when next token will be available)
  const resetIn = allowed ? 0 : (1 - bucket.tokens) / refillRate * 1000;

  return {
    allowed,
    remainingTokens: Math.floor(bucket.tokens),
    resetIn: Math.ceil(resetIn),
  };
}

/**
 * Reset rate limit for a user (admin only)
 * 
 * @param userId - User identifier
 */
export function resetRateLimit(userId: string): void {
  const keysToDelete = Array.from(rateLimitBuckets.keys()).filter((key) =>
    key.startsWith(`${userId}:`)
  );

  keysToDelete.forEach((key) => rateLimitBuckets.delete(key));
  console.log(`Rate limit reset for user: ${userId}`);
}

/**
 * Get current rate limit status for a user
 * 
 * @param userId - User identifier
 * @returns Current bucket status
 */
export function getRateLimitStatus(userId: string): {
  [key: string]: { tokens: number; lastRefillTime: number };
} {
  const status: { [key: string]: { tokens: number; lastRefillTime: number } } =
    {};

  rateLimitBuckets.forEach((bucket, key) => {
    if (key.startsWith(`${userId}:`)) {
      status[key] = {
        tokens: bucket.tokens,
        lastRefillTime: bucket.lastRefillTime,
      };
    }
  });

  return status;
}