package com.duiba.tuia.youtui.web.tool;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 访问限制
 *
 * @author wubo
 */
public class QpsLimit {
    private static Logger log = LoggerFactory.getLogger(QpsLimit.class);
    private LoadingCache<String, AtomicInteger> qpsCache;
    private LoadingCache<String, AtomicInteger> qpmCache;

    /**
     * singleton
     */
    private static QpsLimit instance = new QpsLimit();
    public static QpsLimit builder() {
        return instance;
    }

    /**
     * 根据实际使用场景配置
     * initialCapacity,
     * maximumSize
     *
     */
    private QpsLimit() {
        /*
         * 使用场景：
         * 1.猛犸消息，qps限制50,即单台qps=50/6台机器 = 8,缓存量：1
         */
        qpsCache = CacheBuilder
                .newBuilder()
                .refreshAfterWrite(1,TimeUnit.SECONDS)
                .concurrencyLevel(4)
                .initialCapacity(100)
                .maximumSize(10000L)
                .build(new CacheLoader<String, AtomicInteger>() {
                    public AtomicInteger load(String key) throws Exception {
                        return new AtomicInteger(0);
                    }
                });

        /*
         * 使用场景:
         * 1.直投页用户维度防刷限流 ，每个用户qpm限制18，单台qpm=18/6 = 3
         * (直投页每分钟uv 大概2k左右，9.18晚十点 异常高峰 达到5k,缓存量：估计2k)
         */
        qpmCache = CacheBuilder
                .newBuilder()
                .refreshAfterWrite(1,TimeUnit.MINUTES)
                .concurrencyLevel(4)
                .initialCapacity(3000)
                .maximumSize(10000L)
                .build(new CacheLoader<String, AtomicInteger>() {
                    public AtomicInteger load(String key) throws Exception {
                        return new AtomicInteger(0);
                    }
                });

    }

    /**
     * query per second limit
     * @param bizKey key
     * @param limitTimes limit times
     * @param isIncrement 是否增加qps的统计
     * @return boolean
     */
    public boolean isQpsLimit(String bizKey, int limitTimes,boolean isIncrement) {
        return isIncrement ? isLimitAndIncrement(bizKey, limitTimes, qpsCache) : isLimit(bizKey, limitTimes, qpsCache);
    }

    /**
     * query per second increment
     * @param bizKey key
     */
    public void incrementQps(String bizKey) {
        increment(bizKey, qpsCache);
    }

    /**
     * query per minute limit
     * @param bizKey key
     * @param limitTimes limit times
     * @return boolean
     */
    public boolean isQpmLimit(String bizKey, int limitTimes) {
        return isLimitAndIncrement(bizKey, limitTimes, qpmCache);
    }

    /**
     * 判断是否到达限制次数，同时统计加1
     * @param key 缓存key
     * @param limitTimes 限制次数
     * @param cache 缓存
     * @return boolean
     */
    private boolean isLimitAndIncrement(String key, int limitTimes, LoadingCache<String, AtomicInteger> cache) {
        try {
            AtomicInteger times = cache.get(key);
            return times.incrementAndGet() <= limitTimes;
        } catch (ExecutionException e) {
            log.warn("", e);
            return false;
        }
    }

    /**
     * 判断是否到达限制次数
     * @param key 缓存key
     * @param limitTimes 限制次数
     * @param cache 缓存
     * @return boolean
     */
    private boolean isLimit(String key, int limitTimes, LoadingCache<String, AtomicInteger> cache) {
        try {
            AtomicInteger times = cache.get(key);
            return times.get() <= limitTimes;
        } catch (ExecutionException e) {
            log.warn("", e);
            return false;
        }

    }

    /**
     * 统计加1
     * @param key 缓存key
     * @param cache 缓存
     */
    private void increment(String key, LoadingCache<String, AtomicInteger> cache) {
        try {
            AtomicInteger times = cache.get(key);
            times.incrementAndGet();
        } catch (ExecutionException e) {
            log.warn("", e);
        }
    }
}
