package com.pharmacopoeia.security;

import com.pharmacopoeia.config.RateLimitProperties;
import jakarta.servlet.FilterChain;
import org.jetbrains.annotations.NotNull;
import org.slf4j.MDC;
import com.pharmacopoeia.util.IpUtils;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.filter.OncePerRequestFilter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 用户侧限流：仅限 chat POST 接口，每 IP 每分钟限制次数。
 * 超限返回 429，不会进入 Controller 层，因此不会写缓存和 DB。
 */
@Component
public class RateLimitFilter extends OncePerRequestFilter {

    private final RateLimitProperties props;
    private final Map<String, AtomicInteger> minuteCounter = new ConcurrentHashMap<>();

    public RateLimitFilter(RateLimitProperties props) {
        this.props = props;
    }

    @Override
    protected void doFilterInternal(@NotNull HttpServletRequest request,
                                    @NotNull HttpServletResponse response,
                                    @NotNull FilterChain filterChain) throws ServletException, IOException {
        String ip = IpUtils.getClientIp(request);
        MDC.put("ip", ip);
        try {
            boolean isChatApi = request.getRequestURI().contains("/chat/")
                    && "POST".equalsIgnoreCase(request.getMethod());
            if (isChatApi) {
                var ipCount = minuteCounter.computeIfAbsent(ip, k -> new AtomicInteger(0));
                int count = ipCount.incrementAndGet();
                if (count > props.getPerMinute()) {
                    ipCount.decrementAndGet(); // 超限不计
                    writeRateLimitError(response);
                    return;
                }
            }

            filterChain.doFilter(request, response);
        } finally {
            MDC.clear();
        }
    }

    private static void writeRateLimitError(HttpServletResponse response) throws IOException {
        response.setStatus(429);
        response.setContentType("application/json;charset=UTF-8");
        response.getWriter().write(
                "{\"error\":\"" + "请求过于频繁，请稍后再试" + "\",\"detail\":\"" + "请求过于频繁，请稍后再试" + "\",\"code\":429}"
        );
    }
}
