package com.ruoyi.common.filter;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.ruoyi.common.constant.EncryptConstants;
import com.ruoyi.common.constant.SM4Constants;
import com.ruoyi.common.utils.EncryptHttpServletResponseWrapper;
import com.ruoyi.common.utils.SM4Utils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.util.AntPathMatcher;
import org.springframework.util.PathMatcher;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

public class EncryptResponseFilter implements Filter {

    private static final ObjectMapper objectMapper = new ObjectMapper();
    private final PathMatcher pathMatcher = new AntPathMatcher();

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {

        HttpServletRequest httpRequest = (HttpServletRequest) request;
        HttpServletResponse httpResponse = (HttpServletResponse) response;

        // 检查是否开启加解密功能
        if (!EncryptConstants.ENCRYPT_ENABLED) {
            chain.doFilter(request, response);
            return;
        }

        // 检查是否是SSE流式响应 - 不处理
        if (isSseResponse(httpRequest, httpResponse)) {
            chain.doFilter(request, response);
            return;
        }

        // 检查URL是否需要加密响应
        if (!needProcess(httpRequest)) {
            chain.doFilter(request, response);
            return;
        }

        // 包装响应
        EncryptHttpServletResponseWrapper responseWrapper = new EncryptHttpServletResponseWrapper(httpResponse);
        chain.doFilter(request, responseWrapper);

        // 获取原始响应内容
        byte[] content = responseWrapper.getContent();
        String originalResponse = new String(content, StandardCharsets.UTF_8);

        // 加密响应内容
        if (StringUtils.isNotBlank(originalResponse) && !originalResponse.trim().isEmpty()) {
            String encryptedResponse = SM4Utils.encryptEcb(SM4Constants.SM4_KET, originalResponse);

            // 构建加密响应
            Map<String, Object> result = new HashMap<>();
            result.put("encrypted", true);
            result.put("encryptedData", encryptedResponse);
            result.put("timestamp", System.currentTimeMillis());

            String finalResponse = objectMapper.writeValueAsString(result);

            // 设置响应
            byte[] encryptedBytes = finalResponse.getBytes(StandardCharsets.UTF_8);
            httpResponse.setContentLength(encryptedBytes.length);
            httpResponse.setContentType("application/json;charset=UTF-8");
            httpResponse.getOutputStream().write(encryptedBytes);
        } else {
            // 空响应直接返回空
            httpResponse.getOutputStream().write(new byte[0]);
        }
    }

    /**
     * 判断请求是否需要处理
     */
    private boolean needProcess(HttpServletRequest request) {
        String requestURI = request.getRequestURI();

        // 检查URL是否匹配需要加解密的模式
        for (String pattern : EncryptConstants.URL_PATTERNS) {
            if (pathMatcher.match(pattern, requestURI)) {
                return true;
            }
        }

        return false;
    }

    /**
     * 判断是否是SSE流式响应
     */
    private boolean isSseResponse(HttpServletRequest request, HttpServletResponse response) {
        String path = request.getRequestURI();
        String method = request.getMethod();

        // 检查是否是聊天接口
        if ("/app/chat/chat".equals(path) && "POST".equalsIgnoreCase(method)) {
            return true;
        }

        // 检查Accept头是否包含text/event-stream
        String acceptHeader = request.getHeader("Accept");
        if (acceptHeader != null && acceptHeader.contains("text/event-stream")) {
            return true;
        }

        // 检查响应头是否已经设置为text/event-stream
        String contentType = response.getContentType();
        if (contentType != null && contentType.contains("text/event-stream")) {
            return true;
        }

        return false;
    }
}