const Redis = require('ioredis');

class TokenBucketRateLimiter {
  constructor(redis, options = {}) {
    this.redis = redis;
    this.prefix = options.prefix || 'ratelimit:';
    this.defaultBucketSize = options.bucketSize || 20;
    this.defaultRefillRate = options.refillRate || 10; // tokens/sec
  }

  async consume(key, cost = 1, bucketSize = this.defaultBucketSize, refillRate = this.defaultRefillRate) {
    const redisKey = `${this.prefix}${key}`;
    const now = Date.now();

    // Lua script ensures atomicity - prevents race conditions
    const luaScript = `
      local key = KEYS[1]
      local now = tonumber(ARGV[1])
      local cost = tonumber(ARGV[2])
      local bucket_size = tonumber(ARGV[3])
      local refill_rate = tonumber(ARGV[4])

      -- Fetch current state or initialize
      local tokens = tonumber(redis.call('GET', key) or bucket_size)
      local last_refill = tonumber(redis.call('GET', key .. ':last') or now)

      -- Calculate tokens to add based on elapsed time
      local time_passed = (now - last_refill) / 1000
      local tokens_to_add = time_passed * refill_rate
      local new_tokens = math.min(bucket_size, tokens + tokens_to_add)

      -- Attempt to consume tokens
      if new_tokens >= cost then
        new_tokens = new_tokens - cost
        redis.call('SET', key, new_tokens)
        redis.call('SET', key .. ':last', now)
        redis.call('EXPIRE', key, 3600) -- Auto-cleanup after 1 hour
        redis.call('EXPIRE', key .. ':last', 3600)

        return {1, math.floor(new_tokens), 0, math.floor((bucket_size - new_tokens) / refill_rate)}
      else
        -- Not enough tokens - calculate wait time
        local tokens_needed = cost - new_tokens
        local wait_time = tokens_needed / refill_rate
        return {0, 0, math.ceil(wait_time), 0}
      end
    `;

    const result = await this.redis.eval(luaScript, 1, redisKey, now, cost, bucketSize, refillRate);

    return {
      allowed: Boolean(result[0]),
      remaining: Number(result[1]),
      retryAfter: Number(result[2]),
      resetAfter: Number(result[3])
    };
  }

  // Convenience methods for common time units
  async allowPerSecond(key, tokensPerSecond, cost = 1) {
    return this.consume(key, cost, tokensPerSecond * 2, tokensPerSecond);
  }

  async allowPerMinute(key, tokensPerMinute, cost = 1) {
    const rate = tokensPerMinute / 60;
    return this.consume(key, cost, tokensPerMinute, rate);
  }

  async allowPerHour(key, tokensPerHour, cost = 1) {
    const rate = tokensPerHour / 3600;
    return this.consume(key, cost, tokensPerHour, rate);
  }
}

module.exports = TokenBucketRateLimiter;

const rateLimitMiddleware = (limiter, options = {}) => {
  return async (req, res, next) => {
    // Identify user by IP, API key, or custom logic
    const identifier = options.identifier 
      ? options.identifier(req) 
      : req.ip || req.headers['x-forwarded-for'];

    const cost = options.cost || 1;
    const result = await limiter.allowPerSecond(
      identifier, 
      options.tokensPerSecond || 10, 
      cost
    );

    // Set IETF standard rate limit headers
    res.set({
      'RateLimit-Limit': options.bucketSize || 20,
      'RateLimit-Remaining': result.remaining,
      'RateLimit-Reset': Math.ceil(Date.now() / 1000 + result.resetAfter),
      // Legacy X- prefixed headers for compatibility
      'X-RateLimit-Limit': options.bucketSize || 20,
      'X-RateLimit-Remaining': result.remaining,
      'X-RateLimit-Reset': Math.ceil(Date.now() / 1000 + result.resetAfter)
    });

    if (!result.allowed) {
      res.set('Retry-After', result.retryAfter);
      return res.status(429).json({
        error: 'Rate limit exceeded',
        retryAfter: result.retryAfter,
        message: `Too many requests. Retry after ${result.retryAfter} seconds.`
      });
    }

    next();
  };
};

module.exports = rateLimitMiddleware;

const express = require('express');
const Redis = require('ioredis');
const TokenBucketRateLimiter = require('./token-bucket-limiter');
const rateLimitMiddleware = require('./rate-limit-middleware');

const redis = new Redis({ host: 'localhost', port: 6379 });
const limiter = new TokenBucketRateLimiter(redis, {
  prefix: 'api:global:',
  bucketSize: 100,
  refillRate: 20 // 20 req/sec sustained, 100 burst
});

const app = express();

// Apply globally
app.use(rateLimitMiddleware(limiter, {
  tokensPerSecond: 20,
  bucketSize: 100,
  identifier: req => req.headers['x-api-key'] || req.ip
}));

app.get('/api/users', (req, res) => {
  res.json({ users: [] });
});

app.listen(3000);

// Cheap read: 1 token
app.get('/api/posts', 
  rateLimitMiddleware(limiter, { 
    tokensPerSecond: 20, 
    cost: 1 
  }), 
  (req, res) => {
    res.json({ posts: [] });
  }
);

// Expensive search: 5 tokens
app.post('/api/search', 
  rateLimitMiddleware(limiter, { 
    tokensPerSecond: 20, 
    cost: 5 
  }), 
  async (req, res) => {
    const results = await expensiveSearch(req.body.query);
    res.json(results);
  }
);

// Very expensive AI generation: 20 tokens (entire bucket)
app.post('/api/generate', 
  rateLimitMiddleware(limiter, { 
    tokensPerSecond: 20, 
    cost: 20 
  }), 
  async (req, res) => {
    const generated = await aiGenerate(req.body.prompt);
    res.json(generated);
  }
);

const freeLimiter = new TokenBucketRateLimiter(redis, {
  prefix: 'tier:free:',
  bucketSize: 10,
  refillRate: 2 // 2/sec
});

const paidLimiter = new TokenBucketRateLimiter(redis, {
  prefix: 'tier:paid:',
  bucketSize: 100,
  refillRate: 50 // 50/sec
});

app.use((req, res, next) => {
  const user = getUserFromToken(req.headers.authorization);
  const limiter = user.tier === 'paid' ? paidLimiter : freeLimiter;
  const tokensPerSecond = user.tier === 'paid' ? 50 : 2;

  return rateLimitMiddleware(limiter, { 
    tokensPerSecond,
    identifier: () => user.id 
  })(req, res, next);
});

class SlidingWindowRateLimiter {
  async isAllowed(key, limit, windowSec) {
    const now = Date.now();
    const windowStart = now - (windowSec * 1000);

    const luaScript = `
      local key = KEYS[1]
      local now = tonumber(ARGV[1])
      local window_start = tonumber(ARGV[2])
      local limit = tonumber(ARGV[3])

      -- Remove old entries
      redis.call('ZREMRANGEBYSCORE', key, 0, window_start)

      -- Count current entries
      local current = redis.call('ZCARD', key)

      if current < limit then
        redis.call('ZADD', key, now, now)
        redis.call('EXPIRE', key, ${windowSec})
        return 1
      else
        return 0
      end
    `;

    const result = await this.redis.eval(luaScript, 1, key, now, windowStart, limit);
    return Boolean(result);
  }
}

const Prometheus = require('prom-client');

const rateLimitCounter = new Prometheus.Counter({
  name: 'rate_limit_hits_total',
  help: 'Total rate limit checks',
  labelNames: ['allowed', 'route', 'tier']
});

const rateLimitMiddlewareWithMetrics = (limiter, options = {}) => {
  return async (req, res, next) => {
    const result = await limiter.consume(/* ... */);

    rateLimitCounter.inc({
      allowed: result.allowed ? 'yes' : 'no',
      route: req.path,
      tier: req.user?.tier || 'anonymous'
    });

    if (!result.allowed) {
      return res.status(429).json({ error: 'Rate limit exceeded' });
    }

    next();
  };
};

# Rate limit rejection rate
rate(rate_limit_hits_total{allowed="no"}[5m]) 
/ 
rate(rate_limit_hits_total[5m])

# Top rejected routes
topk(5, sum by (route) (rate(rate_limit_hits_total{allowed="no"}[1h])))

const now = await redis.time(); // Returns [seconds, microseconds]
const timestampMs = now[0] * 1000 + Math.floor(now[1] / 1000);

async consume(key, cost) {
  try {
    const result = await this.redis.eval(/* ... */);
    return result;
  } catch (err) {
    console.error('Redis error:', err);
    // Fail open (allow request) or fail closed (deny request)?
    // Production: fail open with logging, alert if error rate > 1%
    return { allowed: true, remaining: 0, retryAfter: 0 };
  }
}

const identifier = req => {
  if (req.headers['x-api-key']) {
    return `apikey:${req.headers['x-api-key']}`;
  }
  return `ip:${req.ip}`;
};

// Limit by API key (100/sec) AND IP (1000/sec)
app.use(rateLimitMiddleware(keyLimiter, { 
  tokensPerSecond: 100,
  identifier 
}));

app.use(rateLimitMiddleware(ipLimiter, { 
  tokensPerSecond: 1000,
  identifier: req => req.ip 
}));

// Before (express-rate-limit)
const rateLimit = require('express-rate-limit');
const limiter = rateLimit({
  windowMs: 60 * 1000,
  max: 100
});
app.use(limiter);

// After (token bucket)
const redis = new Redis();
const limiter = new TokenBucketRateLimiter(redis, {
  bucketSize: 100,
  refillRate: 100 / 60 // 100 per minute = 1.67/sec
});
app.use(rateLimitMiddleware(limiter, {
  tokensPerSecond: 1.67,
  bucketSize: 100
}));

const axios = require('axios');

async function loadTest() {
  const results = { allowed: 0, denied: 0 };

  for (let i = 0; i < 100; i++) {
    try {
      await axios.get('http://localhost:3000/api/test');
      results.allowed++;
    } catch (err) {
      if (err.response?.status === 429) {
        results.denied++;
      }
    }
  }

  console.log(results); // { allowed: 20, denied: 80 } for 20-token bucket
}

loadTest();

// Service A
const limiterA = new TokenBucketRateLimiter(sharedRedis, {
  prefix: 'service-a:',
  bucketSize: 100,
  refillRate: 20
});

// Service B
const limiterB = new TokenBucketRateLimiter(sharedRedis, {
  prefix: 'service-b:',
  bucketSize: 50,
  refillRate: 10
});

const globalLimiter = new TokenBucketRateLimiter(sharedRedis, {
  prefix: 'global:',
  bucketSize: 1000,
  refillRate: 200
});

// Apply both global and service-specific limits
app.use(rateLimitMiddleware(globalLimiter, { /* ... */ }));
app.use(rateLimitMiddleware(limiterA, { /* ... */ }));

app.use(async (req, res, next) => {
  const user = await getUserFromDB(req.userId);
  const config = {
    tokensPerSecond: user.tier === 'premium' ? 100 : 10,
    bucketSize: user.tier === 'premium' ? 500 : 50
  };

  return rateLimitMiddleware(limiter, config)(req, res, next);
});

const limiters = {
  free: new TokenBucketRateLimiter(redis, { bucketSize: 50, refillRate: 10 }),
  premium: new TokenBucketRateLimiter(redis, { bucketSize: 500, refillRate: 100 })
};

app.use((req, res, next) => {
  const tier = req.user?.tier || 'free';
  return rateLimitMiddleware(limiters[tier], { /* ... */ })(req, res, next);
});