backend/utils/rate_limit.ts

import { loadConfig } from "../types/config";
import type { RateLimitInfo } from "../types/models";

const config = loadConfig();

interface RateLimitBucket {
  tokens: number;
  lastRefill: number;
}

class RateLimiter {
  private buckets = new Map<string, RateLimitBucket>();
  private readonly refillInterval = 60000;

  checkRateLimit(apiKey: string, tier: 'default' | 'premium' | 'admin'): RateLimitInfo {
    const limit = this.getLimitForTier(tier);
    const now = Date.now();
    
    let bucket = this.buckets.get(apiKey);
    
    if (!bucket) {
      bucket = {
        tokens: limit,
        lastRefill: now,
      };
      this.buckets.set(apiKey, bucket);
    }

    const timeSinceRefill = now - bucket.lastRefill;
    if (timeSinceRefill >= this.refillInterval) {
      bucket.tokens = limit;
      bucket.lastRefill = now;
    }

    if (bucket.tokens <= 0) {
      const resetAt = bucket.lastRefill + this.refillInterval;
      throw {
        statusCode: 429,
        message: 'Rate limit exceeded',
        limit,
        remaining: 0,
        resetAt,
      };
    }

    bucket.tokens -= 1;

    const resetAt = bucket.lastRefill + this.refillInterval;

    return {
      limit,
      remaining: bucket.tokens,
      reset_at: resetAt,
      tier,
    };
  }

  private getLimitForTier(tier: 'default' | 'premium' | 'admin'): number {
    switch (tier) {
      case 'admin':
        return config.rateLimit.admin;
      case 'premium':
        return config.rateLimit.premium;
      default:
        return config.rateLimit.default;
    }
  }

  getRateLimitInfo(apiKey: string, tier: 'default' | 'premium' | 'admin'): RateLimitInfo {
    const limit = this.getLimitForTier(tier);
    const bucket = this.buckets.get(apiKey);

    if (!bucket) {
      return {
        limit,
        remaining: limit,
        reset_at: Date.now() + this.refillInterval,
        tier,
      };
    }

    return {
      limit,
      remaining: bucket.tokens,
      reset_at: bucket.lastRefill + this.refillInterval,
      tier,
    };
  }

  cleanup(): void {
    const now = Date.now();
    const maxAge = this.refillInterval * 2;

    for (const [key, bucket] of this.buckets.entries()) {
      if (now - bucket.lastRefill > maxAge) {
        this.buckets.delete(key);
      }
    }
  }
}

export const rateLimiter = new RateLimiter();

setInterval(() => {
  rateLimiter.cleanup();
}, 300000);

export function checkRateLimit(apiKey: string, tier: 'default' | 'premium' | 'admin'): RateLimitInfo {
  return rateLimiter.checkRateLimit(apiKey, tier);
}

export function getRateLimitInfo(apiKey: string, tier: 'default' | 'premium' | 'admin'): RateLimitInfo {
  return rateLimiter.getRateLimitInfo(apiKey, tier);
}