package com.bonus.gateway.filter;

import com.bonus.common.core.utils.encryption.Sm4Utils;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;

/**
 * 对返回的 data数据进行加密
 *
 * @author 黑子
 */
@Configuration
@Slf4j
public class ResponseEncryptFilter implements GlobalFilter, Ordered {
  /*  *//**
     * 返回数据是否加密
     *//*
    @Value("${system.encryptResponse}")
    private boolean encryptResponse;*/
    /**
     * 加密标识
     */
    private static final String ENCRYPT_RESPONSE = "encryptResponse";

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        HttpHeaders headers = request.getHeaders();

        // 检查请求头中是否包含加密标志，并且系统是否启用了加密功能
        if (shouldEncrypt(headers)) {
            ServerHttpResponse originalResponse = exchange.getResponse();
            DataBufferFactory bufferFactory = originalResponse.bufferFactory();
            // 设置响应头
            addResponseHeaders(originalResponse, headers);
            // 创建自定义的响应装饰器，用于处理响应数据
            ServerHttpResponseDecorator responseDecorator = buildResponse(originalResponse, bufferFactory);
            // 使用装饰器包装后的响应继续处理链
            return chain.filter(exchange.mutate().response(responseDecorator).build());
        }

        // 如果不需要加密，直接继续处理链
        return chain.filter(exchange);
    }

    @Override
    public int getOrder() {
        return -5;
    }

    /**
     * 判断是否需要对响应数据进行加密
     *
     * @param headers 请求头
     * @return 如果需要加密返回true，否则返回false
     */
    private boolean shouldEncrypt(HttpHeaders headers) {
        return Boolean.parseBoolean(headers.getFirst(ENCRYPT_RESPONSE));
    }

    /**
     * 添加自定义的响应头
     *
     * @param response       原始响应对象
     * @param requestHeaders 请求头
     */
    private void addResponseHeaders(ServerHttpResponse response, HttpHeaders requestHeaders) {
        response.getHeaders().setContentType(MediaType.APPLICATION_JSON);
        response.getHeaders().add(ENCRYPT_RESPONSE, requestHeaders.getFirst(ENCRYPT_RESPONSE));
    }

    /**
     * 创建自定义的响应装饰器，用于加密响应数据
     *
     * @param originalResponse 原始响应对象
     * @param bufferFactory    数据缓冲区工厂
     * @return 自定义的响应装饰器
     */
    private ServerHttpResponseDecorator buildResponse(ServerHttpResponse originalResponse, DataBufferFactory bufferFactory) {
        return new ServerHttpResponseDecorator(originalResponse) {
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                // 只处理状态码为200 OK的响应，并且响应体是Flux类型
                if (getStatusCode().equals(HttpStatus.OK) && body instanceof Flux) {
                    Flux<? extends DataBuffer> fluxBody = Flux.from(body);

                    // 对响应体进行加密处理
                    return super.writeWith(fluxBody.buffer().map(dataBuffers -> {
                        DataBuffer joinedBuffer = joinDataBuffers(dataBuffers);
                        byte[] content = readContent(joinedBuffer);
                        DataBufferUtils.release(joinedBuffer);

                        // 将响应数据加密
                        String responseData = new String(content, StandardCharsets.UTF_8);
                        responseData = Sm4Utils.encrypt(responseData);
                        byte[] encryptedContent = responseData.getBytes(StandardCharsets.UTF_8);
                        // 设置加密后的内容长度
                        originalResponse.getHeaders().setContentLength(encryptedContent.length);
                        return bufferFactory.wrap(encryptedContent);
                    }));
                } else {
                    log.error("Failed to retrieve response body. Status code: {}", getStatusCode());
                }
                return super.writeWith(body);
            }

            @Override
            public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
                return writeWith(Flux.from(body).flatMapSequential(p -> p));
            }
        };
    }

    /**
     * 将多个DataBuffer连接成一个DataBuffer
     *
     * @param dataBuffers DataBuffer列表
     * @return 连接后的DataBuffer
     */
    private DataBuffer joinDataBuffers(java.util.List<? extends DataBuffer> dataBuffers) {
        if (dataBuffers.size() > 1) {
            return new DefaultDataBufferFactory().join(dataBuffers);
        } else {
            return dataBuffers.get(0);
        }
    }

    /**
     * 从DataBuffer中读取字节数组内容
     *
     * @param dataBuffer DataBuffer对象
     * @return 读取到的字节数组
     */
    private byte[] readContent(DataBuffer dataBuffer) {
        byte[] content = new byte[dataBuffer.readableByteCount()];
        dataBuffer.read(content);
        return content;
    }
}

