package com.xjrsoft.common.aspect;

import cn.hutool.core.collection.ListUtil;
import cn.hutool.extra.servlet.ServletUtil;
import com.baomidou.mybatisplus.core.toolkit.StringPool;
import com.xjrsoft.common.annotation.Limit;
import com.xjrsoft.common.exception.MyException;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.context.annotation.Configuration;
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.data.redis.serializer.GenericToStringSerializer;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

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

@Slf4j
@Aspect
@Configuration
public class LimitAspect {

    private final RedisTemplate<String, Object> redisTemplate;

    public LimitAspect(RedisTemplate<String, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Around("@annotation(com.xjrsoft.common.annotation.Limit)")
    public Object around(ProceedingJoinPoint pjp){
        MethodSignature methodSignature = (MethodSignature)pjp.getSignature();
        Method method = methodSignature.getMethod();
        Limit annotation = method.getAnnotation(Limit.class);

        String key = annotation.key();
        HttpServletRequest request = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();
        String allKey = key + StringPool.COLON + ServletUtil.getClientIP(request);


        int time = annotation.time();
        int count = annotation.count();

        try {
            String luaScript = buildLuaScript();
            RedisScript<Long> redisScript = new DefaultRedisScript<>(luaScript, Long.class);
            // execute使用的redis的默认的序列化方式，需要设置参数--arg的序列化方式，以及result结果的序列化方式
            // 此处传参只要能转为Object就行(因为数字不能直接强转为String，所以不能用String序列化)
            // 结果的类型需要根据脚本定义，此处是数字--定义的是Long类型
            Long result = redisTemplate.execute(redisScript, new GenericToStringSerializer<>(Object.class),
                    new GenericToStringSerializer<>(Long.class), ListUtil.toList(allKey), time, count);

            // 使用lua脚本执行，保证原子性,进行请求次数限制
            assert result != null;
            boolean isOutOfLimit = result.equals(1L);
            if(!isOutOfLimit){
                throw new MyException(annotation.msg());
            }
            else {
               return pjp.proceed();
            }

        }catch (Throwable e){
            if(e instanceof MyException){
                log.debug("令牌桶={}，获取令牌失败",allKey);
                throw new MyException(e.getLocalizedMessage());
            }
            e.printStackTrace();
            throw new MyException("服务器异常");
        }
    }

    public String buildLuaScript(){

        //lua 脚本，进行请求次数的叠加，并判断请求次数是否超过限制
        return "local val = redis.call('incr', KEYS[1]) " +
                "local expire = tonumber(ARGV[1]) " +
                "if val == 1 " +
                "then redis.call('expire', KEYS[1], expire) " +
                "else if redis.call('ttl', KEYS[1]) == -1 " +
                "then redis.call('expire', KEYS[1], expire) " +
                "end " +
                "end " +
                "if val > tonumber(ARGV[2]) " +
                "then return 0 " +
                "end " +
                "return 1";
    }

}