package com.npf.imgserver.aspect;

import com.google.common.util.concurrent.RateLimiter;
import com.npf.imgserver.common.base.ResultBase;
import com.npf.imgserver.common.enums.ResultCodeEnum;
import com.npf.imgserver.utils.JsonUtil;
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.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Scope;
import org.springframework.stereotype.Component;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;

@Component
@Scope
@Aspect
public class RateLimitAop {

    public static final Logger logger = LoggerFactory.getLogger(RateLimitAop.class);

    @Autowired
    private HttpServletResponse response;

    //秒并发限制 50/s
    private RateLimiter rateLimiter = RateLimiter.create(50.0);

    @Pointcut("@annotation(com.npf.imgserver.aspect.RateLimitAspect)")
    public void serviceLimit() {

    }

    @Around("serviceLimit()")
    public Object around(ProceedingJoinPoint joinPoint) {
        Boolean flag = rateLimiter.tryAcquire();
        Object obj = null;
        try {
            if (flag) {
                obj = joinPoint.proceed();
            } else {
                ResultBase resultBase = new ResultBase(ResultCodeEnum.RATE_LIMIT_WARN);
                String result = JsonUtil.object2json(resultBase);
                output(response, result);
            }
        } catch (Throwable e) {
            logger.error("限流信息输出异常", e);
        }
        return obj;
    }

    /**
     * 限流信息输出
     *
     * @param response
     * @param msg
     * @throws IOException
     */
    public void output(HttpServletResponse response, String msg) throws IOException {
        response.setContentType("application/json;charset=UTF-8");
        ServletOutputStream outputStream = null;
        try {
            outputStream = response.getOutputStream();
            outputStream.write(msg.getBytes("UTF-8"));
        } catch (IOException e) {
            logger.error("限流信息IO异常", e);
        } finally {
            outputStream.flush();
            outputStream.close();
        }
    }
}