package com.kxs.gateway.api.filter;

import cn.hutool.crypto.symmetric.AES;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.kxs.common.core.constant.enums.ErrorTypeEnum;
import com.kxs.common.core.exception.AesDecodeException;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
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 org.springframework.web.util.UriComponentsBuilder;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.util.concurrent.atomic.AtomicInteger;



/**
 * 请求解密处理
 *
 * @author 没秃顶的码农
 * @date 2023/11/17
 */
@Slf4j
@RequiredArgsConstructor
public class RequestDecoderFilter extends AbstractGatewayFilterFactory<Object> {

	private final AES cryptoAes;

	private static final String PARAM_NAME = "value";

	@Override
	public GatewayFilter apply(Object config) {
		return (exchange, chain) -> {
			ServerHttpRequest request = exchange.getRequest();
			// 请求方法
			HttpMethod method = request.getMethod();

			if (method == HttpMethod.GET) {
				// 1 修改请求参数,并获取请求参数
				return decryptGetRequestParam(exchange, chain);
			}

			if (method == HttpMethod.POST) {
				return decryptPostBodyParam(exchange, chain);
			}

			return chain.filter(exchange);

		};

	}

	/**
	 * 解密post-body参数
	 *
	 * @param exchange exchange
	 * @param chain    chain
	 * @return mono＜void＞
	 */
	private Mono<Void> decryptPostBodyParam(ServerWebExchange exchange, GatewayFilterChain chain) {
		//获取请求体,修改请求体
		ServerRequest serverRequest = ServerRequest.create(exchange,
				HandlerStrategies.withDefaults().messageReaders());

		//请求长度
		AtomicInteger setContentLength = new AtomicInteger();
		//获取请求头
		MediaType contentType = exchange.getRequest().getHeaders().getContentType();
		Flux<DataBuffer> modifiedBody = serverRequest.bodyToFlux(String.class).flatMap(body -> {
			// 解密请求体
			try {
				log.info("待解密请求参数：{}", body);
				String decryptStr = cryptoAes.decryptStr(body);
				log.info("解密后参数：{}", decryptStr);

				DataBufferFactory dataBufferFactory = exchange.getResponse().bufferFactory();
				DataBuffer bodyDataBuffer = dataBufferFactory.wrap(decryptStr.getBytes());
				//获取参数长度
//					setContentLength.set(encrypt.getBytes().length);
				//添加到请求体
				return Flux.just(bodyDataBuffer);
			} catch (Exception e) {
				log.error(ErrorTypeEnum.DECRYPT_ERROR.getDescription(), e);
				return Flux.error(new AesDecodeException(ErrorTypeEnum.DECRYPT_ERROR.getDescription()));
			}
		});
		ServerHttpRequest mutatedRequest = new ServerHttpRequestDecorator(
				exchange.getRequest()) {
			@NotNull
			@Override
			public HttpHeaders getHeaders() {
				//重新设置请求头
				HttpHeaders httpHeaders = new HttpHeaders();
				httpHeaders.putAll(super.getHeaders());
				httpHeaders.remove(HttpHeaders.CONTENT_TYPE);
				httpHeaders.remove(HttpHeaders.CONTENT_LENGTH);
				assert contentType != null;
				httpHeaders.set(HttpHeaders.CONTENT_TYPE, contentType.toString());

				return httpHeaders;
			}

			@NotNull
			@Override
			public Flux<DataBuffer> getBody() {
				return modifiedBody;
			}
		};
		return chain.filter(exchange.mutate().request(mutatedRequest).build());
	}


	/**
	 * 解密GET请求参数
	 *
	 * @param exchange exchange
	 * @param chain    chain
	 * @return mono＜void＞
	 */
	private Mono<Void> decryptGetRequestParam(ServerWebExchange exchange, GatewayFilterChain chain) {

		try {
			MultiValueMap<String, String> queryParams = exchange.getRequest().getQueryParams();
			String decryptedParam = queryParams.getFirst(PARAM_NAME);
			log.info("GET请求待解密数据 :{}", decryptedParam);
			String decryptStr = cryptoAes.decryptStr(decryptedParam);
			log.info("GET解密后数据 :{}", decryptStr);

			JSONObject decryptJson = JSON.parseObject(decryptStr);
			//构建新的请求参数webflux
			MultiValueMap<String, String> parameters = new LinkedMultiValueMap<>(decryptJson.size());
			decryptJson.forEach((key, value) -> {
				parameters.add(key, (String) value);
			});
			// 构建修改后的URI
			URI modifiedUri = UriComponentsBuilder.fromUri(exchange.getRequest().getURI())
					.replaceQueryParams(parameters)
					.build(true)
					.toUri();
			// 修改请求，设置新的URI
			ServerHttpRequest modifiedRequest = exchange.getRequest().mutate()
					.uri(modifiedUri)
					.build();
			// 用修改后的请求继续处理请求链
			return chain.filter(exchange.mutate().request(modifiedRequest).build());
		} catch (Exception e) {
			log.error(ErrorTypeEnum.DECRYPT_ERROR.getDescription(), e);
			return Mono.error(new AesDecodeException(ErrorTypeEnum.DECRYPT_ERROR.getDescription()));
		}

	}

}
