package com.huimv.management.limit.aspect;

import cn.hutool.core.date.DateTime;
import cn.hutool.core.date.DateUnit;
import cn.hutool.core.date.DateUtil;
import com.google.common.collect.ImmutableList;
import com.huimv.common.exception.LimitAccessException;
import com.huimv.common.utils.HttpContextUtil;
import com.huimv.common.utils.IPUtil;
import com.huimv.management.limit.enums.LimitType;
import com.huimv.management.limit.annotation.Limit;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;



/**
 * 接口限流
 *
 * @author yinhao
 */
@Slf4j
@Aspect
@Component
public class LimitAspect extends AspectSupport {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    @Pointcut("@annotation(com.huimv.management.limit.annotation.Limit)")
    public void pointcut() {
    }

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        // 获取到 HttpServletRequest
        HttpServletRequest request = HttpContextUtil.getHttpServletRequest();
        Method method = resolveMethod(point);
        // 获取到注解
        Limit limitAnnotation = method.getAnnotation(Limit.class);
        // CUSTOMER 还是 IP
        LimitType limitType = limitAnnotation.limitType();
        // 在 redis 中需要使用这个 name 拼接 key
        String name = limitAnnotation.name();
        String key;
        // 获取客户端ip
        String ip = IPUtil.getIpAddr(request);

        //long limitPeriod = limitAnnotation.period();
        //修改过期时间为下一天00:00
        DateTime current = DateUtil.date();
        DateTime tomorrow = DateUtil.beginOfDay(DateUtil.offsetDay(current, 1));
        long expire = DateUtil.between(current, tomorrow, DateUnit.SECOND);


        long limitCount = limitAnnotation.count();
        switch (limitType) {
            case IP:
                key = limitAnnotation.key() + ip;
                break;
            case CUSTOMER:
                key = limitAnnotation.key();
                break;
            default:
                key = StringUtils.upperCase(method.getName());
        }
        // redis 通过key来区分唯一
        ImmutableList<String> keys = ImmutableList.of(StringUtils.join(limitAnnotation.prefix() + "_", key));
        String luaScript = buildLuaScript();
        RedisScript<Long> redisScript = new DefaultRedisScript<>(luaScript, Long.class);
        Long count = redisTemplate.execute(redisScript, keys, limitCount, expire);
        log.info("IP:{} 第 {} 次访问key为 {}，描述为 [{}] 的接口", ip, count, keys, name);
        if (count != null && count.intValue() <= limitCount) {
            return point.proceed();
        } else {
            throw new LimitAccessException("当天可访问次数已达上限!");
        }
    }

    /**
     * 限流脚本 参考redis文档  http://doc.redisfans.com/script/eval.html
     * 调用的时候不超过阈值，则直接返回并执行计算器自加。
     *
     * @return lua脚本
     */
    private String buildLuaScript() {
        return "local c" +
                "\n c = redis.call('get',KEYS[1])" +
                "\n if c and tonumber(c) > tonumber(ARGV[1]) then" +
                "\n return c;" +
                "\n end" +
                "\n c = redis.call('incr',KEYS[1])" +
                "\n if tonumber(c) == 1 then" +
                "\n redis.call('expire',KEYS[1],ARGV[2])" +
                "\n end" +
                "\n return c;";
    }
}
