package org.jeecg.common.limit;


import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import org.apache.poi.ss.formula.functions.T;
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.aspectj.lang.reflect.MethodSignature;
import org.jeecg.common.api.vo.Result;
import org.jeecg.common.util.RedisUtil;
import org.springframework.aop.aspectj.MethodInvocationProceedingJoinPoint;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.lang.reflect.Method;
import java.nio.charset.StandardCharsets;

/**
 * 切面类：实现限流校验
 * @author niugeng
 * @date 2023-12-13 14:02
 */
@Slf4j
@Aspect
@Component
public class AccessLimitAspect {

    @Resource
    private RedisUtil sbRedisUtil;

    /**
     * 未知
     */
    private static final String UNKNOWN = "unknown";

    /**
     * 这里我们使用注解的形式
     * 当然，我们也可以通过切点表达式直接指定需要拦截的package,需要拦截的class 以及 method
     */
    @Pointcut("@annotation(org.jeecg.common.limit.RateLimit)")
    public void limitPointCut() {
    }

    /**
     * 环绕通知
     */
    @Around("limitPointCut()")
    public Object around(ProceedingJoinPoint pjp) throws Throwable {
        // 获取被注解的方法
        MethodInvocationProceedingJoinPoint mjp = (MethodInvocationProceedingJoinPoint) pjp;
        MethodSignature signature = (MethodSignature) mjp.getSignature();
        Method method = signature.getMethod();

        // 获取方法上的注解
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);
        if (rateLimit == null) {
            // 如果没有注解，则继续调用，不做任何处理
            return pjp.proceed();
        }
        /**
         * 代码走到这里，说明有 RateLimit 注解，那么就需要做限流校验了
         *  使用Redis的API做计数校验
         */
        //获取request对象
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if (attributes == null) {
            return pjp.proceed();
        }
        HttpServletRequest request = attributes.getRequest();
        // 获取请求IP地址
        String ip = getIpAddr(request);
        // 请求url路径
        String uri = request.getRequestURI();
        //存到redis中的key
        String key = rateLimit.key() + ip + ":" + uri;
        // 缓存中存在key，在限定访问周期内已经调用过当前接口
        if (sbRedisUtil.hasKey(key)) {
            // 访问次数自增1
            sbRedisUtil.incr(key, 1);
            // 超出访问次数限制
            if (Integer.parseInt(StrUtil.toString(sbRedisUtil.get(key))) > rateLimit.count()) {

                Result<T> fail = Result.error(rateLimit.msg());

                String utf8Str = StrUtil.utf8Str(JSON.toJSONString(fail));

                renderString(attributes.getResponse(), utf8Str);

                //访问限制，不继续执行原方法
                return utf8Str;
            }
            // 未超出访问次数限制，不进行任何操作，返回true
        } else {
            // 第一次设置数据,过期时间为注解确定的访问周期
            sbRedisUtil.set(key, 1, rateLimit.cycle());
        }
        return pjp.proceed();
    }


    /**
     * 获取请求的归属IP地址
     *
     * @param request 请求
     * @return {@link String}
     */
    private String getIpAddr(HttpServletRequest request) {
        String ipAddress = null;
        try {
            ipAddress = request.getHeader("x-forwarded-for");
            if (ipAddress == null || ipAddress.length() == 0 || UNKNOWN.equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getHeader("Proxy-Client-IP");
            }
            if (ipAddress == null || ipAddress.length() == 0 || UNKNOWN.equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getHeader("WL-Proxy-Client-IP");
            }
            if (ipAddress == null || ipAddress.length() == 0 || UNKNOWN.equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getRemoteAddr();
            }
            // 对于通过多个代理的情况，第一个IP为客户端真实IP,多个IP按照','分割
            if (ipAddress != null && ipAddress.length() > 15 && ipAddress.indexOf(",") >=0) {
                ipAddress = ipAddress.substring(0, ipAddress.indexOf(","));
            }
        } catch (Exception e) {
            ipAddress = "";
        }
        return ipAddress;
    }

    /**
     * 将字符串渲染到客户端
     *
     * @param response 渲染对象
     * @param string   待渲染的字符串
     * @return null
     */
    public static void renderString(HttpServletResponse response, String string) {
        try {
            byte[] bytes = string.getBytes(StandardCharsets.UTF_8);
            response.setHeader("Content-length", String.valueOf(bytes.length));
            response.setStatus(200);
            response.setContentType("application/json");
            response.setCharacterEncoding("utf-8");
            response.getOutputStream().write(bytes);
        } catch (IOException e) {
            log.error("限流切面报错；", e);
        }
    }

}
