import { RateLimiterMemory, RateLimiterRedis } from 'rate-limiter-flexible';
import { redisService } from '../services/redis.service';

interface RateLimitOptions {
  points: number;        // 限流点数
  duration: number;      // 时间窗口（秒）
  blockDuration?: number; // 封锁时间（秒）
  keyGenerator?: (ctx: any) => string; // 自定义键生成器
}

// 内存限流器（Redis 不可用时使用）
const memoryLimiters = new Map<string, RateLimiterMemory>();

// Redis 限流器
const redisLimiters = new Map<string, RateLimiterRedis>();

/**
 * 获取或创建限流器
 */
function getLimiter(name: string, options: RateLimitOptions) {
  if (redisService.isAvailable()) {
    if (!redisLimiters.has(name)) {
      const limiter = new RateLimiterRedis({
        storeClient: redisService['client'],
        keyPrefix: `rate_limit:${name}:`,
        points: options.points,
        duration: options.duration,
        blockDuration: options.blockDuration || options.duration * 2,
      });
      redisLimiters.set(name, limiter);
    }
    return redisLimiters.get(name)!;
  } else {
    if (!memoryLimiters.has(name)) {
      const limiter = new RateLimiterMemory({
        points: options.points,
        duration: options.duration,
      });
      memoryLimiters.set(name, limiter);
    }
    return memoryLimiters.get(name)!;
  }
}

/**
 * 创建限流中间件
 * @param options 限流配置
 */
export function createRateLimiter(options: RateLimitOptions) {
  const limiter = getLimiter('default', options);

  return async (ctx: any, next: any) => {
    try {
      const key = options.keyGenerator 
        ? options.keyGenerator(ctx) 
        : ctx.ip;

      await limiter.consume(key, 1);
      return await next();
    } catch (rejRes: any) {
      const retrySecs = Math.round(rejRes.msBeforeNext / 1000) || 1;
      
      ctx.set('Retry-After', String(retrySecs));
      ctx.status = 429;
      ctx.body = {
        code: 429,
        message: '请求过于频繁，请稍后再试',
        retryAfter: retrySecs,
      };
    }
  };
}

/**
 * API 全局限流（每分钟 100 次）
 */
export const apiRateLimiter = createRateLimiter({
  points: 100,
  duration: 60,
  keyGenerator: (ctx) => `api:${ctx.ip}`,
});

/**
 * 登录接口限流（每分钟 5 次）
 */
export const loginRateLimiter = createRateLimiter({
  points: 5,
  duration: 60,
  blockDuration: 300, // 封锁 5 分钟
  keyGenerator: (ctx) => `login:${ctx.ip}`,
});

/**
 * 验证码发送限流（每分钟 1 次，每小时 5 次）
 */
export const smsRateLimiter = createRateLimiter({
  points: 1,
  duration: 60,
  blockDuration: 3600,
  keyGenerator: (ctx) => `sms:${ctx.request.body?.phone || ctx.ip}`,
});

/**
 * TTS 生成限流（根据用户等级）
 */
export const ttsRateLimiter = createRateLimiter({
  points: 20,
  duration: 60,
  keyGenerator: (ctx) => `tts:${ctx.state.user?.id || ctx.ip}`,
});

/**
 * 文件上传限流（每分钟 10 次）
 */
export const uploadRateLimiter = createRateLimiter({
  points: 10,
  duration: 60,
  keyGenerator: (ctx) => `upload:${ctx.state.user?.id || ctx.ip}`,
});
