package org.springblade.common.interceptor;

import org.springblade.common.annotation.AccessLimit;
import org.springblade.common.utils.IPUtils;
import org.springblade.core.tool.api.R;
import org.springblade.core.tool.jackson.JsonUtil;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.ValueOperations;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;
import org.springframework.web.servlet.ModelAndView;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.lang.reflect.Method;
import java.util.concurrent.TimeUnit;

/**
 * 接口限流，防盗刷拦截器
 */
@Component
public class AccessLimitInterceptor implements HandlerInterceptor  {

    @Autowired
    private StringRedisTemplate redisTemplate;


    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (!(handler instanceof HandlerMethod)) {
            return true;
        }

        HandlerMethod handlerMethod = (HandlerMethod) handler;
        Method method = handlerMethod.getMethod();
        AccessLimit accessLimit = method.getAnnotation(AccessLimit.class);
        if (accessLimit == null) {
            return true;
        }

        int second = accessLimit.second();
        int maxCount = accessLimit.maxCount();
        String uri = request.getRequestURI();
        String ip = IPUtils.getIpAddr(request);
        String key = ip + ":" + uri;

        ValueOperations<String, String> ops1 = redisTemplate.opsForValue();
        Integer count = null;
        boolean flag = redisTemplate.hasKey(key);
        if(flag){
            count = Integer.valueOf(ops1.get(key));
        }
        //第一次访问记录到redis，第二次则累加1，否则超出访问次数
        if (null == count) {
            ops1.set(key, "1", second, TimeUnit.SECONDS);
        } else if (count < maxCount) {
            ops1.increment(key, 1);
        } else {
            //超出访问次数。这里可以记录该ip，登录信息，请求地址等用于后期跟踪
            return setResponse(JsonUtil.toJson(R.fail("请求过于频繁")), response);
        }

        return true;
    }

    private boolean setResponse(String message, HttpServletResponse response) throws IOException {
        ServletOutputStream outputStream = null;
        try {
            response.setHeader("Content-type", "application/json; charset=utf-8");
            outputStream = response.getOutputStream();
            outputStream.write(message.getBytes("UTF-8"));
        } catch (Exception ex) {
            return false;
        } finally {
            if (outputStream != null) {
                outputStream.flush();
                outputStream.close();
            }
        }
        return false;
    }

    @Override
    public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) throws Exception {

    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {

    }
}
