package com.chushang.gateway.filter;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.IdUtil;
import com.chushang.gateway.config.GatewayContext;
import io.netty.buffer.ByteBufAllocator;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.IOUtils;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.ResolvableType;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.HttpMethod;
import org.springframework.http.InvalidMediaTypeException;
import org.springframework.http.MediaType;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.codec.ServerCodecConfigurer;
import org.springframework.http.codec.multipart.Part;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Objects;

/**
 * by zhaowenyuan create 2022/6/8 14:27
 * 通用请求过滤
 */
@Component
@Slf4j
public class RequestGlobalFilter implements GlobalFilter, Ordered
{

    /**
     * default HttpMessageReader
     */
    private static final List<HttpMessageReader<?>> messageReaders = HandlerStrategies.withDefaults().messageReaders();

    private static final ResolvableType MULTIPART_DATA_TYPE = ResolvableType.forClassWithGenerics(MultiValueMap.class, String.class, Part.class);

    private static final Mono<MultiValueMap<String, Part>> EMPTY_MULTIPART_DATA = Mono.just(CollectionUtils.unmodifiableMultiValueMap(new LinkedMultiValueMap<String, Part>(0))).cache();

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        if(request.getMethod() == HttpMethod.GET){
            // get请求 处理参数
            return handleGetMethod(exchange, chain, request);
        }

        if(request.getMethod() == HttpMethod.POST){
            // post请求 处理参数
            return handlePostMethod(exchange, chain, request);
        }

        return chain.filter(exchange);
    }

    /**
     * get请求 处理参数
     */
    private Mono<Void> handleGetMethod(ServerWebExchange exchange, GatewayFilterChain chain, ServerHttpRequest request) {

        return chain.filter(exchange);
    }

    private Mono<Void> handlePostMethod(ServerWebExchange exchange, GatewayFilterChain chain, ServerHttpRequest request){
        GatewayContext gatewayContext = new GatewayContext();
        gatewayContext.setPath(request.getPath().pathWithinApplication().value());
        //save gateway context into exchange
        exchange.getAttributes().put(GatewayContext.CACHE_GATEWAY_CONTEXT, gatewayContext);

        MediaType contentType = request.getHeaders().getContentType();
        if(MediaType.APPLICATION_JSON.equals(contentType)
                || MediaType.APPLICATION_JSON_UTF8.equals(contentType)){
            // 请求内容为 application json

            // 重新构造 请求体
            return readJsonBody(exchange, chain, gatewayContext);
        }
//        else
//        if (MediaType.MULTIPART_FORM_DATA.isCompatibleWith(contentType)) {
//            // 请求内容为 form data
//            return readFormData(exchange, chain, gatewayContext);
//        }
        return chain.filter(exchange);
    }

    private Mono<Void> readJsonBody(ServerWebExchange exchange, GatewayFilterChain chain, GatewayContext gatewayContext) {
        return DataBufferUtils.join(exchange.getRequest().getBody())
                .flatMap(dataBuffer -> {
                    /*
                     * read the body Flux<DataBuffer>, and release the buffer
                     * //TODO when SpringCloudGateway Version Release To G.SR2,this can be update with the new version's feature
                     * see PR https://github.com/spring-cloud/spring-cloud-gateway/pull/1095
                     */
                    byte[] bytes = new byte[dataBuffer.readableByteCount()];
                    dataBuffer.read(bytes);
                    DataBufferUtils.release(dataBuffer);
                    @SuppressWarnings("ReactiveStreamsUnusedPublisher") Flux<DataBuffer> cachedFlux = Flux.defer(() -> {
                        DataBuffer buffer = exchange.getResponse().bufferFactory().wrap(bytes);
                        DataBufferUtils.retain(buffer);
                        return Mono.just(buffer);
                    });
                    ServerHttpRequest mutatedRequest =
                            new ServerHttpRequestDecorator(exchange.getRequest()) {
                                @SuppressWarnings("NullableProblems")
                                @Override
                                public Flux<DataBuffer> getBody() {
                                    return cachedFlux;
                                }
                            };
                    ServerWebExchange mutatedExchange = exchange.mutate().request(mutatedRequest).build();
                    // save body into gatewayContext
                    return ServerRequest.create(mutatedExchange, messageReaders)
                            .bodyToMono(String.class)
                            .doOnNext(gatewayContext::setCacheBody)
                            .then(chain.filter(mutatedExchange));
                });
    }

    private Mono<Void> readFormData(ServerWebExchange exchange, GatewayFilterChain chain, GatewayContext gatewayContext) {
        return exchange.getRequest().getBody().collectList().flatMap(dataBuffers -> {
            @SuppressWarnings("OptionalGetWithoutIsPresent") final byte[] totalBytes = dataBuffers.stream().map(dataBuffer -> {
                try {
                    return IOUtils.toByteArray(dataBuffer.asInputStream());
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }).reduce(this::addBytes).get();
            final ServerHttpRequestDecorator decorator = new ServerHttpRequestDecorator(exchange.getRequest()) {
                @SuppressWarnings("NullableProblems")
                @Override
                public Flux<DataBuffer> getBody() {
                    return Flux.just(buffer(totalBytes));
                }
            };
            final ServerCodecConfigurer configurer = ServerCodecConfigurer.create();
            final Mono<MultiValueMap<String, Part>> multiValueMapMono = repackageMultipartData(decorator, configurer);
            return multiValueMapMono.publishOn(Schedulers.boundedElastic()).publishOn(Schedulers.boundedElastic())
                    .publishOn(Schedulers.boundedElastic()).flatMap(part -> {
                for (String key : part.keySet()) {
                    // 如果为文件时 则进入下一次循环
                    if (key.equals("file")) {
                        continue;
                    }
                    //noinspection CallingSubscribeInNonBlockingScope
                    Objects.requireNonNull(part.getFirst(key)).content().subscribe(buffer -> {
                        final byte[] bytes = new byte[buffer.readableByteCount()];
                        buffer.read(bytes);
                        DataBufferUtils.release(buffer);
                        final String bodyString = new String(bytes, StandardCharsets.UTF_8);
                        gatewayContext.setCacheBody(bodyString);
                    });
                }
                return chain.filter(exchange.mutate().request(decorator).build());
            });
        });
    }

    @SuppressWarnings("unchecked")
    private static Mono<MultiValueMap<String, Part>> repackageMultipartData(ServerHttpRequest request, ServerCodecConfigurer configurer) {
        try {
            final MediaType contentType = request.getHeaders().getContentType();
            if (MediaType.MULTIPART_FORM_DATA.isCompatibleWith(contentType)) {
                return ((HttpMessageReader<MultiValueMap<String, Part>>) configurer.getReaders().stream().filter(reader -> reader.canRead(MULTIPART_DATA_TYPE, MediaType.MULTIPART_FORM_DATA))
                        .findFirst().orElseThrow(() -> new IllegalStateException("No multipart HttpMessageReader."))).readMono(MULTIPART_DATA_TYPE, request, Collections.emptyMap())
                        .switchIfEmpty(EMPTY_MULTIPART_DATA).cache();
            }
        } catch (InvalidMediaTypeException ex) {
            // Ignore
        }
        return EMPTY_MULTIPART_DATA;
    }

    /**
     * addBytes.
     * @param first first
     * @param second second
     * @return byte
     */
    public byte[] addBytes(byte[] first, byte[] second) {
        final byte[] result = Arrays.copyOf(first, first.length + second.length);
        System.arraycopy(second, 0, result, first.length, second.length);
        return result;
    }

    private DataBuffer buffer(byte[] bytes) {
        final NettyDataBufferFactory nettyDataBufferFactory = new NettyDataBufferFactory(ByteBufAllocator.DEFAULT);
        final DataBuffer buffer = nettyDataBufferFactory.allocateBuffer(bytes.length);
        buffer.write(bytes);
        return buffer;
    }


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