Kordant/apps/api/src/middleware/rate-limit.middleware.ts

import { FastifyInstance, FastifyRequest, FastifyReply } from 'fastify';
import { apiEnv, rateLimitConfig } from '../config/api.config';

// Simple in-memory rate limiter
// In production, this should use Redis or similar distributed store
class RateLimiter {
  private store: Map<string, { count: number; resetTime: number }>;

  constructor() {
    this.store = new Map();
  }

  async checkLimit(
    key: string,
    windowMs: number,
    maxRequests: number
  ): Promise<{ remaining: number; resetTime: number; retryAfter?: number }> {
    const now = Date.now();
    const current = this.store.get(key);

    if (!current || now > current.resetTime) {
      // Reset window
      this.store.set(key, {
        count: 1,
        resetTime: now + windowMs,
      });

      return {
        remaining: maxRequests - 1,
        resetTime: now + windowMs,
      };
    }

    // Increment counter
    current.count++;
    this.store.set(key, current);

    const remaining = maxRequests - current.count;

    if (current.count > maxRequests) {
      return {
        remaining: 0,
        resetTime: current.resetTime,
        retryAfter: current.resetTime - now,
      };
    }

    return {
      remaining,
      resetTime: current.resetTime,
    };
  }

  reset(key: string) {
    this.store.delete(key);
  }
}

const rateLimiter = new RateLimiter();

export async function rateLimitMiddleware(fastify: FastifyInstance) {
  fastify.addHook('preHandler', async (request: FastifyRequest, reply: FastifyReply) => {
    // Skip rate limiting for health checks
    if (request.url === '/health') {
      return;
    }

    // Get client identifier (IP or API key)
    const clientIp = request.ip || request.headers['x-forwarded-for'] || 'unknown';
    const apiKey = request.headers['x-api-key'] as string | undefined;
    const key = apiKey ? `api:${apiKey}` : `ip:${clientIp}`;

    // Determine tier based on API key or default to basic
    let tier = 'basic';
    if (apiKey) {
      // In production, fetch tier from user/service lookup
      // For now, use a simple heuristic based on key format
      if (apiKey.startsWith('premium_')) {
        tier = 'premium';
      } else if (apiKey.startsWith('plus_')) {
        tier = 'plus';
      }
    }

    const config = rateLimitConfig[tier as keyof typeof rateLimitConfig];
    const result = await rateLimiter.checkLimit(
      key,
      config.windowMs,
      config.maxRequests
    );

    // Set rate limit headers
    reply.header('X-RateLimit-Limit', config.maxRequests);
    reply.header('X-RateLimit-Remaining', result.remaining);
    reply.header('X-RateLimit-Reset', Math.ceil(result.resetTime / 1000));

    if (result.retryAfter) {
      reply.header('Retry-After', Math.ceil(result.retryAfter / 1000));
      reply.code(429); // Too Many Requests

      return {
        error: 'Too Many Requests',
        message: `Rate limit exceeded. Try again in ${Math.ceil(result.retryAfter / 1000)}s`,
        tier,
        limit: config.maxRequests,
        reset: new Date(result.resetTime).toISOString(),
      };
    }

    // Add tier info to request for downstream use
    (request as any).rateLimitTier = tier;
  });
}

// Export for testing
export { rateLimiter };