package com.qmth.boot.redis.rateLimit;

import com.qmth.boot.core.rateLimit.entity.RateLimitRule;
import com.qmth.boot.core.rateLimit.entity.RateLimitScope;
import com.qmth.boot.core.rateLimit.service.RateLimitService;
import com.qmth.boot.core.rateLimit.service.RateLimiter;
import com.qmth.boot.core.uid.service.UidMachineService;
import com.qmth.boot.redis.constant.RedisConstant;
import org.redisson.api.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.util.concurrent.TimeUnit;

@Component("rateLimitService")
public class RedisRateLimitService implements RateLimitService, RedisConstant {

    private static final Logger log = LoggerFactory.getLogger(RedisRateLimitService.class);

    private RedissonClient redissonClient;

    private UidMachineService uidMachineService;

    public RedisRateLimitService(RedissonClient redissonClient, UidMachineService uidMachineService) {
        this.redissonClient = redissonClient;
        this.uidMachineService = uidMachineService;
        log.info("Redis rate limit service inited.");
    }

    @Override
    public RateLimiter getRateLimiter(String endpoint, String target, RateLimitRule rule) {
        String key = RATE_LIMIT_PREFIX.concat(KEY_SEPARATOR).concat(endpoint);
        //控制范围为非集群时，加入machineId作为本实例标识
        if (rule.getScope() == RateLimitScope.INSTANCE) {
            key = key.concat(KEY_SEPARATOR).concat(String.valueOf(uidMachineService.getMachineId()));
        }
        // 若包含限流对象，则key加入对象标识
        if (target != null) {
            key = key.concat(KEY_SEPARATOR).concat(target);
        }
        //设置了限流时间段，使用Redisson内置限流器
        if (rule.getPeriod().toMillis() > 0) {
            RRateLimiter limiter = redissonClient.getRateLimiter(key);
            //Key已存在但配置不同需要先删除
            if (limiter.isExists() && !compareConfig(limiter.getConfig(), rule)) {
                limiter.delete();
            }
            limiter.trySetRate(RateType.OVERALL, rule.getCount(), rule.getPeriod().toMillis(),
                    RateIntervalUnit.MILLISECONDS);
            //限流单元默认24小时后自动清理
            limiter.expireAsync(24, TimeUnit.HOURS);
            return new RedisRateLimiter(limiter);
        }
        //限流时间段为0时，使用Redisson内置信号量控制全局并发数
        else {
            RSemaphore semaphore = redissonClient.getSemaphore(key);
            semaphore.trySetPermits(rule.getCount());
            //限流单元默认24小时后自动清理
            semaphore.expireAsync(24, TimeUnit.HOURS);
            return new RedisRateLimiter(semaphore);
        }
    }

    private boolean compareConfig(RateLimiterConfig config, RateLimitRule rule) {
        return config != null && config.getRateType().equals(RateType.OVERALL) && config.getRate() == rule.getCount()
                && config.getRateInterval() == rule.getPeriod().toMillis();
    }

}
