package com.idmt.difyweb.filter;

import com.idmt.difyweb.common.WebUtil;
import jakarta.servlet.*;
import jakarta.servlet.annotation.WebFilter;
import jakarta.servlet.http.*;
import org.springframework.core.annotation.Order;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.net.http.HttpResponse;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

@FunctionalInterface
interface MyLogger {
    void apply(String message);
}

@Component
public class RateLimitFilter implements Filter {
    private static final Logger logger = LoggerFactory.getLogger(RateLimitFilter.class);
    MyLogger myLogger=logger::info; //System.out::println; //
    // 配置参数
    //统计一分钟的限流 1分钟60秒 1秒1000毫秒
    private static  final int MAX_Restricted_Milliseconds = 10000; //10 * 1000; //10秒
    private static final int MAX_REQUESTS_PER_MINUTE = 10; //150的请求，指定时间内
    private static final int MAX_404_COUNT = 10; //5次返回 404，指定时间内
    private static final long BLOCK_TIME_MILLIS = 30 * 60 * 1000; // 限制 30分钟
    private static final long CLEANUP_INTERVAL = 2 * 60 * 1000; // 5分钟清理一次
    private static class IpStats {
        private final AtomicInteger requestCount = new AtomicInteger(0);//0;
        private final AtomicInteger error404Count =  new AtomicInteger(0); //0;
        private final AtomicLong lastResetTime = new AtomicLong(System.currentTimeMillis());

        public boolean isExpired(long currentTime) {
            return currentTime - lastResetTime.get() > MAX_Restricted_Milliseconds;
        }
        public void reset(long currentTime) {
            requestCount.set(0);
            error404Count.set(0);
            lastResetTime.set(currentTime);
        }
    }

    // 存储IP的请求统计信息
    private final ConcurrentHashMap<String, IpStats> ipStatsMap = new ConcurrentHashMap<>();
    // 存储被封禁的IP及其解封时间
    private final ConcurrentHashMap<String, Long> blackListMap = new ConcurrentHashMap<>();

    private ScheduledExecutorService cleanupExecutor;

    private void printRresult(ConcurrentHashMap<String, Long> hashMap){
        for (Map.Entry<?, ?> entry : hashMap.entrySet()) {
            myLogger.apply("Key: " + entry.getKey() + ", Value: " + entry.getValue());
        }
    }
    private void printRresultIPStats(ConcurrentHashMap<String, IpStats> hashMap){
        for (Map.Entry<?, ?> entry : hashMap.entrySet()) {
            IpStats ipStats=(IpStats)entry.getValue();
            myLogger.apply("Key: " + entry.getKey() +
                    ", requestCount: " + ipStats.requestCount.get()+
                    ", error404Count: " + ipStats.error404Count.get()+
                    ", lastResetTime:"+ipStats.lastResetTime.get()
                    );
        }
    }
    private void printALl(){
        myLogger.apply("---封闭的IP----");
        printRresult(blackListMap);
        myLogger.apply("---IP的404次数----");
        printRresultIPStats(ipStatsMap);
    }
    //封闭IP
    private void blockIP(String ip,long currentTime){
        myLogger.apply("开始封闭IP");
        blackListMap.put(ip, currentTime + BLOCK_TIME_MILLIS);
        // 可选: 清理统计数据
        ipStatsMap.remove(ip);
    }
    /**
     * 清理过期数据，防止内存泄漏
     */
    private void cleanup() {
        myLogger.apply("execute cleanup");
        long currentTime = System.currentTimeMillis();
        // 清理过期IP统计
        ipStatsMap.entrySet().removeIf(entry ->
                entry.getValue().isExpired(currentTime)
        );
        // 清理已解封的IP
        blackListMap.entrySet().removeIf(entry ->
                entry.getValue() <= currentTime
        );
        printALl();
    }
    @Override
    public void init(FilterConfig filterConfig) {

        // 定期清理过期IP和黑名单
        cleanupExecutor = Executors.newSingleThreadScheduledExecutor(r -> {
            Thread t = new Thread(r, "RateLimit-Cleanup");
            t.setDaemon(true);
            return t;
        });
        //cleanupExecutor.scheduleAtFixedRate(this::cleanup, CLEANUP_INTERVAL, CLEANUP_INTERVAL, TimeUnit.MILLISECONDS);
        cleanupExecutor.scheduleWithFixedDelay(this::cleanup, CLEANUP_INTERVAL, CLEANUP_INTERVAL, TimeUnit.MILLISECONDS);
    }
    @Override
    public void destroy() {
        myLogger.apply("--开始清理--");
        if (cleanupExecutor != null) {
            cleanupExecutor.shutdown();
        }
    }
    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
       /* myLogger.apply("--开始过滤--");*/
        printALl();
        HttpServletRequest request1=(HttpServletRequest)request;
        String ip = WebUtil.getClientIp(request1); //request.getRemoteAddr();
        long currentTime = System.currentTimeMillis();

        // 检查是否被封禁
        Long blockUntil = blackListMap.get(ip);
        if (blockUntil != null && blockUntil > currentTime) {
            HttpServletResponse res=(HttpServletResponse) response;
            res.setStatus(429); // 太多请求
            response.setContentType("text/plain;charset=UTF-8");
            response.getWriter().write("Access restricted, please try again later。");
            return;
        } else if (blockUntil != null && blockUntil <= currentTime) {
            // 解封
            blackListMap.remove(ip);
        }
        // 统计请求信息
        IpStats stats = ipStatsMap.computeIfAbsent(ip, k -> new IpStats());

        // 存在跨分钟的情况，重置统计
        if (stats.isExpired(currentTime)) {
            stats.reset(currentTime);
        }
        // 请求计数
        stats.requestCount.incrementAndGet();
        ipStatsMap.put(ip, stats);

        // 3. 检查限流（在 请求处理前 检查）
        if (stats.requestCount.get() > MAX_REQUESTS_PER_MINUTE) {
            blockIP(ip,currentTime);
            HttpServletResponse res=(HttpServletResponse) response;
            res.setStatus(429);
            res.setContentType("text/plain;charset=UTF-8");
            res.getWriter().write("Requests are too frequent, please try again later");
            return;
        }

        // 处理请求
        // 使用自定义Response包装器捕获状态码
        StatusCaptureResponseWrapper wrappedResponse = new StatusCaptureResponseWrapper((HttpServletResponse) response);
        try {
            chain.doFilter(request, wrappedResponse);
        }
        finally {
            // 判断响应状态码
            if (wrappedResponse.getStatus() == 404) {
                stats.error404Count.incrementAndGet();
            }

            // 超过404阈值，加入黑名单
            if (stats.error404Count.get() >= MAX_404_COUNT) {
                blockIP(ip,currentTime);
            }
        }
    }

    /**
     * 自定义ResponseWrapper，用于捕获状态码
     */
    private static class StatusCaptureResponseWrapper extends HttpServletResponseWrapper {
        private int httpStatus = 200;

        public StatusCaptureResponseWrapper(HttpServletResponse response) {
            super(response);
        }

        @Override
        public void setStatus(int sc) {
            super.setStatus(sc);
            this.httpStatus = sc;
        }

        @Override
        public void sendError(int sc) throws IOException {
            super.sendError(sc);
            this.httpStatus = sc;
        }

        @Override
        public void sendError(int sc, String msg) throws IOException {
            super.sendError(sc, msg);
            this.httpStatus = sc;
        }

        public int getStatus() {
            return this.httpStatus;
        }
    }
}