package com.ruoyi.framework.aspect;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.ruoyi.common.enums.RateLimitType;
import com.ruoyi.common.utils.ContextUtil;
import com.ruoyi.common.utils.RedisUtil;
import com.ruoyi.common.utils.SecurityUtil;
import com.ruoyi.common.web.exception.ServiceException;
import com.ruoyi.framework.annotation.RyRateLimit;
import org.noear.solon.annotation.Inject;
import org.noear.solon.core.aspect.Invocation;
import org.noear.solon.core.aspect.MethodInterceptor;
import org.noear.solon.core.wrap.MethodHolder;

import java.util.Collections;

/**
 * 限流处理
 * @author ruoyi
 */
public class RateLimitAspect implements MethodInterceptor {

    @Inject
    RedisUtil redisUtil;

    @Override
    public Object doIntercept(Invocation inv) throws Throwable {
        return handleIntercept(inv);
    }

    /**
     * 限流实现
     * @param inv
     * @return
     */
    private Object handleIntercept(Invocation inv) throws Throwable {
        MethodHolder method = inv.method();
        RyRateLimit ryRateLimit = method.getAnnotation(RyRateLimit.class);
        // 拼接 key，形如：rate_limit:127.0.0.1-com.ruoyi.framework.filter.GlobalFilter-doFilter
        StringBuilder sb = new StringBuilder(ryRateLimit.key());
        if(ryRateLimit.type() == RateLimitType.TOKEN) {
            // 以 token限制
            String token = SecurityUtil.getTokenValue();
            if(StrUtil.isNotBlank(token)) {
                sb.append(token).append("-");
            }
        } else if (ryRateLimit.type() == RateLimitType.IP) {
            // 以 ip限制
            sb.append(ContextUtil.getIpAddr()).append("-");
        }
        sb.append(method.getDeclaringClz().getName()).append("-").append(method.getMethod().getName());
        int count = ryRateLimit.count();
        Long num = redisUtil.executeScript(limitScriptText(), Collections.singletonList(sb.toString()), count, ryRateLimit.time());
        if (ObjectUtil.isNull(num) || num.intValue() > count) {
            throw new ServiceException("访问过于频繁，请稍候再试");
        }
        return inv.invoke();
    }

    /**
     * 限流脚本
     */
    private String limitScriptText() {
        return "local key = KEYS[1]\n" +
                "local count = tonumber(ARGV[1])\n" +
                "local time = tonumber(ARGV[2])\n" +
                "local current = redis.call('get', key);\n" +
                "if current and tonumber(current) > count then\n" +
                "    return tonumber(current);\n" +
                "end\n" +
                "current = redis.call('incr', key)\n" +
                "if tonumber(current) == 1 then\n" +
                "    redis.call('expire', key, time)\n" +
                "end\n" +
                "return tonumber(current);";
    }

}
