package com.mvc.interceptor;

import com.alibaba.fastjson.JSON;
import com.mvc.model.ResultGeekerModel;
import com.mvc.service.aiiden.AiIdenService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpStatus;
import org.springframework.web.servlet.HandlerInterceptor;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.Queue;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;

/**
 * IP限流拦截器
 */
public class IPRateLimitInterceptor implements HandlerInterceptor {
    private static Logger log = LoggerFactory.getLogger(AiIdenService.class);

    /**
     * 在TIME_WINDOW时间内,访问次数超过MAX_REQUESTS后,限制访问 BLOCK_DURATION 时长,后恢复(限制期间内,再访问,恢复时间不变)
     */
    private static final int MAX_REQUESTS = 600;            //次数限制(暂时设置的宽裕,考虑同一个主网下,多个客户端访问的情况)
    private static final long TIME_WINDOW = 60 * 1000;      // 1分钟内
    private static final long BLOCK_DURATION = 60000;       // 1分钟(多久恢复)

    private Map<String, Queue<Long>> requestTimes = new ConcurrentHashMap<>();
    private Map<String, Long> blockedIPs = new ConcurrentHashMap<>();

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        String ipAddress = request.getRemoteAddr();
        //log.info("远程的端口是："+request.getRemotePort()+"Host:"+request.getRemoteHost()+"User:"+request.getRemoteUser());

        // 检查是否被禁止访问
        if (isIPBlocked(ipAddress)) {
            response.setStatus(HttpStatus.FORBIDDEN.value());
            handleErrorResponse(response,"You are blocked from accessing this resource.");
            //response.getWriter().write("You are blocked from accessing this resource.");
            return false;
        }

        // 获取当前时间戳
        long currentTime = System.currentTimeMillis();

        // 检查是否超过限流
        if (isRateLimited(ipAddress, currentTime)) {
            // 添加到被禁止访问的IP列表
            blockIP(ipAddress);
            response.setStatus(HttpStatus.TOO_MANY_REQUESTS.value());
            handleErrorResponse(response,"Too many requests. Please try again later.");
            log.info("对IP[" + ipAddress + "]开始限流");
            return false;
        }

        // 更新请求计数
        updateRequestCount(ipAddress, currentTime);

        return true;
    }

    /**
     * 是否限流
     *
     * @param ipAddress
     * @param currentTime
     * @return
     */
    private boolean isRateLimited(String ipAddress, long currentTime) {
        Queue<Long> times = requestTimes.get(ipAddress);
        if (times == null) {
            return false;
        }
        long timeWindowStart = currentTime - TIME_WINDOW;
        int count = 0;
        for (long time : times) {
            if (time >= timeWindowStart) {
                count++;
            }
        }
        return count >= MAX_REQUESTS;
    }

    /**
     * 更新请求次数
     *
     * @param ipAddress
     * @param currentTime
     */
    private void updateRequestCount(String ipAddress, long currentTime) {
        Queue<Long> times = requestTimes.computeIfAbsent(ipAddress, k -> new ConcurrentLinkedQueue<>());
        times.offer(currentTime);
    }

    /**
     * 判断指定 IP 地址是否被禁止访问
     *
     * @param ipAddress IP 地址
     * @return 是否被禁止访问
     */
    private boolean isIPBlocked(String ipAddress) {
        Long blockedTime = blockedIPs.get(ipAddress);
        return blockedTime != null && System.currentTimeMillis() - blockedTime < BLOCK_DURATION;
    }

    /**
     * 将指定 IP 地址加入禁止访问列表，并设置禁止时间为当前时间，5分钟后自动解除禁止
     *
     * @param ipAddress IP 地址
     */
    private void blockIP(String ipAddress) {
        blockedIPs.put(ipAddress, System.currentTimeMillis());
        // 创建计时器，5分钟后解除IP的禁止访问状态
        Timer timer = new Timer();
        timer.schedule(new TimerTask() {
            @Override
            public void run() {
                unblockIP(ipAddress);
                timer.cancel();
            }
        }, BLOCK_DURATION);
    }

    /**
     * 对IP限制进行解除
     *
     * @param ipAddress
     */
    private void unblockIP(String ipAddress) {
        log.info("对IP[" + ipAddress + "]解除限流");
        blockedIPs.remove(ipAddress);
    }

    /**
     * 处理错误响应
     * @param response
     * @param msg
     */
    private void handleErrorResponse (HttpServletResponse response,String msg) {
        // 返回一个错误响应，例如401 Unauthorized状态码和一个错误消息
        try{
            response.setCharacterEncoding("UTF-8");
            //response.setStatus(1502);
            response.setContentType("application/json");
            //response.sendError(1502);
            ResultGeekerModel model=new ResultGeekerModel(msg);
            model.setData("");
            String json = JSON.toJSONString(model);

            byte[] bytes = json.getBytes(StandardCharsets.UTF_8);
            response.getOutputStream().write(bytes);
            /*//response.getWriter().write(json);
            //使用getWriter()会有此异常:nested exception is java.lang.IllegalStateException: getWriter() has already been called for this response
            //无法同时使用getWriter()和getOutputStream(),怀疑SpringBoot默认已经使用了使用getWriter */
        }catch (Exception ex){
            System.out.println(ex.getMessage());
        }
    }
}