Rate Limiting

import http from "node:http";
import Redis from "ioredis";

const redis = new Redis(process.env.REDIS_URL || "redis://localhost:6379");

// --- Sliding Window Rate Limiter ---
class SlidingWindowRateLimiter {
  constructor(
    private maxRequests: number,
    private windowMs: number
  ) {}

  async isAllowed(key: string): Promise<{
    allowed: boolean;
    remaining: number;
    resetAt: number;
    retryAfter: number;
  }> {
    const now = Date.now();
    const windowStart = now - this.windowMs;
    const redisKey = `rl:${key}`;

    // Use a Redis pipeline for atomicity
    const pipeline = redis.pipeline();

    // Remove expired entries
    pipeline.zremrangebyscore(redisKey, 0, windowStart);

    // Count current entries
    pipeline.zcard(redisKey);

    // Add current request (we'll remove it if denied)
    pipeline.zadd(redisKey, now, `${now}:${Math.random()}`);

    // Set TTL on the key
    pipeline.pexpire(redisKey, this.windowMs);

    const results = await pipeline.exec();
    const currentCount = (results?.[1]?.[1] as number) || 0;

    if (currentCount >= this.maxRequests) {
      // Over limit — remove the entry we just added
      await redis.zremrangebyscore(redisKey, now, now);

      // Find when the oldest entry expires
      const oldest = await redis.zrange(redisKey, 0, 0, "WITHSCORES");
      const oldestTime = oldest.length >= 2 ? parseInt(oldest[1]) : now;
      const retryAfter = Math.ceil((oldestTime + this.windowMs - now) / 1000);

      return {
        allowed: false,
        remaining: 0,
        resetAt: oldestTime + this.windowMs,
        retryAfter: Math.max(retryAfter, 1),
      };
    }

    return {
      allowed: true,
      remaining: this.maxRequests - currentCount - 1,
      resetAt: now + this.windowMs,
      retryAfter: 0,
    };
  }
}

// --- Token Bucket Rate Limiter ---
class TokenBucketRateLimiter {
  constructor(
    private capacity: number,       // max tokens
    private refillRate: number,      // tokens added per second
  ) {}

  async isAllowed(key: string): Promise<{
    allowed: boolean;
    remaining: number;
    retryAfter: number;
  }> {
    const redisKey = `rl:tb:${key}`;
    const now = Date.now();

    // Lua script for atomic token bucket
    const script = `
      local key = KEYS[1]
      local capacity = tonumber(ARGV[1])
      local refill_rate = tonumber(ARGV[2])
      local now = tonumber(ARGV[3])

      local bucket = redis.call('HMGET', key, 'tokens', 'last_refill')
      local tokens = tonumber(bucket[1])
      local last_refill = tonumber(bucket[2])

      if tokens == nil then
        tokens = capacity
        last_refill = now
      end

      -- Refill tokens based on elapsed time
      local elapsed = (now - last_refill) / 1000
      tokens = math.min(capacity, tokens + (elapsed * refill_rate))
      last_refill = now

      local allowed = 0
      if tokens >= 1 then
        tokens = tokens - 1
        allowed = 1
      end

      redis.call('HMSET', key, 'tokens', tokens, 'last_refill', last_refill)
      redis.call('PEXPIRE', key, math.ceil(capacity / refill_rate) * 1000)

      return {allowed, math.floor(tokens)}
    `;

    const result = (await redis.eval(
      script, 1, redisKey,
      this.capacity, this.refillRate, now
    )) as number[];

    const allowed = result[0] === 1;
    const remaining = result[1];

    return {
      allowed,
      remaining,
      retryAfter: allowed ? 0 : Math.ceil(1 / this.refillRate),
    };
  }
}

// --- Rate Limit Middleware ---
type RateLimitTier = {
  maxRequests: number;
  windowMs: number;
};

const tiers: Record<string, RateLimitTier> = {
  free:       { maxRequests: 100, windowMs: 60 * 60 * 1000 },   // 100/hour
  pro:        { maxRequests: 1000, windowMs: 60 * 60 * 1000 },  // 1000/hour
  enterprise: { maxRequests: 10000, windowMs: 60 * 60 * 1000 }, // 10000/hour
};

function getUserTier(_req: http.IncomingMessage): string {
  // In production: look up user's plan from auth token
  return "free";
}

function getClientKey(req: http.IncomingMessage): string {
  // Use API key or IP address
  const apiKey = req.headers["x-api-key"];
  if (apiKey) return `api:${apiKey}`;
  return `ip:${req.socket.remoteAddress}`;
}

async function rateLimitMiddleware(
  req: http.IncomingMessage,
  res: http.ServerResponse
): Promise<boolean> {
  const tier = tiers[getUserTier(req)];
  const limiter = new SlidingWindowRateLimiter(tier.maxRequests, tier.windowMs);
  const key = getClientKey(req);

  const result = await limiter.isAllowed(key);

  // Always set rate limit headers
  res.setHeader("X-RateLimit-Limit", tier.maxRequests);
  res.setHeader("X-RateLimit-Remaining", result.remaining);
  res.setHeader("X-RateLimit-Reset", Math.ceil(result.resetAt / 1000));

  if (!result.allowed) {
    res.setHeader("Retry-After", result.retryAfter);
    res.writeHead(429, { "Content-Type": "application/json" });
    res.end(JSON.stringify({
      error: "Too many requests",
      retryAfter: result.retryAfter,
    }));
    return false;
  }

  return true;
}

// --- Server ---
const server = http.createServer(async (req, res) => {
  const allowed = await rateLimitMiddleware(req, res);
  if (!allowed) return;

  res.writeHead(200, { "Content-Type": "application/json" });
  res.end(JSON.stringify({ message: "OK", timestamp: new Date().toISOString() }));
});

server.listen(3000, () => console.log("Server with rate limiting on :3000"));