package com.fschool.hc.gateway.filter;


import com.fschool.hc.gateway.util.SystemConstants;
import org.springframework.core.annotation.Order;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.List;

@Component
public class CacheReqBodyFilter implements WebFilter {
    private static final List<HttpMessageReader<?>> messageReaders = HandlerStrategies.withDefaults().messageReaders();

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        if(exchange.getRequest().getHeaders().getContentLength() > 0) {
            //Join all the DataBuffers so we have a single DataBuffer for the body
            return DataBufferUtils.join(exchange.getRequest().getBody())
                    .flatMap(dataBuffer -> {
                        //Update the retain counts so we can read the body twice, once to parse into an object
                        //Note: if we end up reading the body twice we will run into a problem
                        //now there is no good use case for doing this
                        DataBufferUtils.retain(dataBuffer);
                        //Make a slice for each read so each read has its own read/write indexes
                        Flux<DataBuffer> cachedFlux = Flux.defer(() -> Flux.just(dataBuffer.slice(0, dataBuffer.readableByteCount())));

                        ServerHttpRequest mutatedRequest = new ServerHttpRequestDecorator(exchange.getRequest()) {
                            @Override
                            public Flux<DataBuffer> getBody() {
                                return cachedFlux;
                            }
                        };

                        ServerWebExchange mutatedExchange = exchange.mutate().request(mutatedRequest).build();

                        return ServerRequest.create(exchange.mutate().request(mutatedRequest).build(), messageReaders)
                                .bodyToMono(String.class)
                                .doOnNext(objectValue -> exchange.getAttributes().put(SystemConstants.CACHED_REQ_BODY, objectValue))
                                .then(chain.filter(mutatedExchange));
                    });
        }
        return chain.filter(exchange);
    }
}