import { Injectable } from '@nestjs/common';
import * as Redis from 'ioredis';
import { ConfigService } from '@app/common';
import { createHash } from 'crypto'
import { v4 as uuidv4 } from 'uuid';

const rateLimitScript = ""http:// 后面單獨(dú)列出

@Injectable()
export class RedisService {
    private readonly redisClient: Redis.Redis;
    private luaScript: any;
    constructor(
        private readonly configService: ConfigService,
    ) {

        const self = this;
        const connConfig = this.configService.get("redisService")
        this.redisClient = new Redis.Redis(connConfig)
        this.luaScript = {
            rateLimit: {
                script: rateLimitScript,
                hash: self.hashStr(rateLimitScript)
            },
        }
        
    }

    private hashStr(value: string) {
        return createHash("sha1").update(value).digest('hex')
    }

    async rateLimit(opts: any): Promise<boolean> {

        const { key, limit, windowSize } = opts;
        const uuid = uuidv4()
        let result;
        const { script, hash } = this.luaScript.rateLimit
  
        try {

            const shaResult = await this.redisClient.evalsha(hash, 1, key, limit, windowSize, uuid)
            result = shaResult

        } catch (error) {

            const shaResult = await this.redisClient.eval(script, 1, key, limit, windowSize, uuid)
            result = shaResult

        }
        return result == 1
    }
}

--傳入四個(gè)參數(shù) 分別是key,限制次數(shù),時(shí)間范圍,唯一值
local key = KEYS[1]
local limit = tonumber(ARGV[1])
local windowSize = tonumber(ARGV[2]) --單位毫秒
local uuid = ARGV[3] -- 唯一值是為了防止zset 重復(fù)

-- 使用redis 來獲取時(shí)間，防止多進(jìn)程生成相似的邊界導(dǎo)致超頻。時(shí)間單位是微秒
local date = redis.call("time")
local now = tonumber(date[1]) * 1000000 + tonumber(date[2])
local startTime = now - windowSize * 1000
local endTime = now +  1000000

-- 計(jì)算過期時(shí)間 時(shí)間單位是秒
local expireSec = tonumber(math.ceil(windowSize / 1000)) + 1

-- 統(tǒng)計(jì)當(dāng)前zset數(shù)組里的數(shù)據(jù),超出范圍則返回0,
-- 否則做3件事,然后返回1
-- 1、向數(shù)組里增加新值
-- 2、刪除數(shù)組中開始時(shí)間之前的數(shù)據(jù),防止數(shù)組過大
-- 3、給數(shù)組續(xù)過期時(shí)間
local count = tonumber(redis.call('zcount', key, startTime, endTime))

if count + 1 > limit then
    return 0
else
    redis.call('zadd', key, now, uuid)
    redis.call('zremrangebyscore', key, 0, startTime - 100000)
    redis.call('expire', key, expireSec)
    return 1
end

import { SetMetadata } from '@nestjs/common';

export interface rateLimitOptions {
    keyPrefix: string,
    limit: number,
    windowSize: number
}

export const RateLimit = (options: rateLimitOptions): MethodDecorator => SetMetadata('rateLimit', options)

import { Injectable, ExecutionContext, CanActivate } from "@nestjs/common";
import { Reflector } from '@nestjs/core';
import { RedisService, rateLimitOptions } from "@app/common";
import { BusinessException } from "@app/common";

@Injectable()
export class RateLimitGuard implements CanActivate {
    constructor(
        private reflector: Reflector,
        private redisService: RedisService
    ) { }

    private getIpFromRequest(request: { ip: string }): string {
        return request.ip?.match(/\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b/)?.[0]
    }

    async canActivate(context: ExecutionContext) {

         // 通過反射拿到前面設(shè)置的值
        const rateLimitConfig = this.reflector.get<rateLimitOptions>("rateLimit", context.getHandler());
           if (!rateLimitConfig) {
			 // 當(dāng)前接口如果沒設(shè)置參數(shù)則定義默認(rèn)參數(shù)
            const cMethod = this.reflector.get("method", context.getHandler());// 是GET,POST 等http method
            const cPath = this.reflector.get("path", context.getHandler());// 接口的具體路徑
            rateLimitConfig = {
                keyPrefix: cMethod + ":" + cPath,
                limit: 1,
                windowSize: 5000
            }
        }

        const { keyPrefix, limit, windowSize } = rateLimitConfig
        const request = context.switchToHttp().getRequest();
        const ip = this.getIpFromRequest(request)
        const key = keyPrefix + ":" + ip

        const isPass = await this.redisService.rateLimit({
            key,
            limit,
            windowSize
        })
        if (!isPass) {
            // 返回自定義的錯(cuò)誤
            throw new BusinessException("RATE_LIMIT_EXCEEDED_LIMIT")
        }
        return true
    }
}

@Controller('user')
export class UserController {
  constructor(private readonly userService: UserService) { }
  
  @Public()
  @RateLimit({ keyPrefix: "login", limit: 3, windowSize: 1000 })
  @UseGuards(RateLimitGuard)
  @Post('register')
  register(@Body() createUserDto: CreateUserDto) {
    return this.userService.register(createUserDto);
   }
}

@Module({
  providers: [
    {
      provide: APP_GUARD,
      useClass: RateLimitGuard
    }
  ],
})