package com.red.circle.gateway.filter;

import com.google.common.base.Throwables;
import com.red.circle.framework.core.request.ServiceTraceId;
import com.red.circle.gateway.logging.GatewayLogging;
import com.red.circle.gateway.logging.GatewayLoggingResponse;
import com.red.circle.gateway.logging.LoggingUploadService;
import com.red.circle.gateway.util.ApiLoggingUtils;
import com.red.circle.tool.core.json.JacksonUtils;
import com.red.circle.tool.core.text.StringUtils;
import lombok.RequiredArgsConstructor;
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.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
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.http.HttpHeaders;
import org.springframework.http.HttpStatusCode;
import org.springframework.http.MediaType;
import org.springframework.http.ReactiveHttpOutputMessage;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
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 java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.zip.GZIPInputStream;

/**
 * 全局拦截器，应用所有的微服务, 对请求的API调用过滤，记录接口的请求时间，方便日志审计、告警、分析等运维操作.
 *
 * @author pengliang on 2023/4/18
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class ApiLoggingFilter implements GlobalFilter, Ordered {

    private static final int MAX_RESPONSE_BODY_LOG_LENGTH = 500;

    private final LoggingUploadService loggingUploadService;
    private final List<HttpMessageReader<?>> messageReaders = HandlerStrategies.withDefaults()
            .messageReaders();

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {

        ServerHttpRequest request = exchange.getRequest();
        String userAgent = request.getHeaders().getFirst("User-Agent");
        if (StringUtils.isNotBlank(userAgent) && userAgent.startsWith("ELB-HealthChecker")) {
            return chain.filter(exchange);
        }

        GatewayLogging logging = ApiLoggingUtils.createGatewayLogging(exchange);

        return checkBodyType(exchange.getRequest())
                ? writeBodyLog(exchange, chain, logging)
                : writeBasicLog(exchange, chain, logging);
    }

    private boolean checkBodyType(ServerHttpRequest request) {
        MediaType mediaType = request.getHeaders().getContentType();
        return MediaType.APPLICATION_FORM_URLENCODED.isCompatibleWith(mediaType)
                || MediaType.APPLICATION_JSON.isCompatibleWith(mediaType);
    }


    private Mono<Void> writeBasicLog(ServerWebExchange exchange, GatewayFilterChain chain,
                                     GatewayLogging logging) {

        // 请求参数
        logging.getRequest().setBody(ApiLoggingUtils.getRequestParamsToString(exchange));
        // 获取响应体
        ServerHttpResponseDecorator decoratedResponse = responseLog(exchange, logging);
        // 包装一个新的req
        ServerHttpRequest request = exchange.getRequest().mutate().headers(this::setRequestTraceId)
                .build();
        return chain.filter(exchange.mutate().request(request).response(decoratedResponse).build())
                .then(Mono.fromRunnable(() -> {
                    ApiLoggingUtils.resetGatewayLogging(exchange, logging);
                    logging.uploadLog(loggingUploadService);
                }));
    }

    /**
     * 解决 request body 只能读取一次问题， 参考:
     * org.springframework.cloud.gateway.filter.factory.rewrite.ModifyRequestBodyGatewayFilterFactory
     */
    private Mono<Void> writeBodyLog(ServerWebExchange exchange, GatewayFilterChain chain,
                                    GatewayLogging logging) {

        ServerRequest serverRequest = ServerRequest.create(exchange, messageReaders);

        Mono<Object> modifiedBody = serverRequest.bodyToMono(Object.class)
                .flatMap(body -> {
                    logging.getRequest().setBody(JacksonUtils.toJson(body));
                    return Mono.just(body);
                });

        // 通过 BodyInserter 插入 body(支持修改body), 避免 request body 只能获取一次
        BodyInserter<Mono<Object>, ReactiveHttpOutputMessage> bodyInserter = BodyInserters.fromPublisher(
                modifiedBody, Object.class);
        HttpHeaders headers = new HttpHeaders();
        headers.putAll(exchange.getRequest().getHeaders());
        // the new content type will be computed by bodyInserter
        // and then set in the request decorator
        headers.remove(HttpHeaders.CONTENT_LENGTH);

        CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, headers);

        return bodyInserter.insert(outputMessage, new BodyInserterContext())
                .then(Mono.defer(() -> {
                    // 重新封装请求
                    ServerHttpRequest decoratedRequest = requestDecorate(exchange, headers, outputMessage);

                    // 记录响应日志
                    ServerHttpResponseDecorator decoratedResponse = responseLog(exchange, logging);

                    // 记录普通的
                    return chain.filter(
                                    exchange.mutate().request(decoratedRequest).response(decoratedResponse).build())
                            .then(Mono.fromRunnable(() -> {
                                ApiLoggingUtils.resetGatewayLogging(exchange, logging);
                                logging.uploadLog(loggingUploadService);
                            }));
                }));
    }


    /**
     * 记录响应日志 通过 DataBufferFactory 解决响应体分段传输问题.
     */
    @SuppressWarnings("all")
    private ServerHttpResponseDecorator responseLog(ServerWebExchange exchange, GatewayLogging gatewayLog) {
        ServerHttpResponse response = exchange.getResponse();
        DataBufferFactory bufferFactory = response.bufferFactory();

        return new ServerHttpResponseDecorator(response) {
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                GatewayLoggingResponse loggingResponse = gatewayLog.getResponse();
                loggingResponse.setTime(System.currentTimeMillis());
                gatewayLog.calculateExecuteTime();
                loggingResponse.setCode(Optional.ofNullable(response.getStatusCode())
                        .map(HttpStatusCode::value)
                        .map(Objects::toString)
                        .orElse("?"));
                if (body instanceof Flux && !isSuccessResponse(response)) {
                    // 获取响应类型，如果是 json 就打印
                    if (isJsonResponse(exchange, response)) {
                        return super.writeWith(DataBufferUtils.join(Flux.from(body)).map(dataBuffer -> {
                            byte[] content = new byte[dataBuffer.readableByteCount()];
                            dataBuffer.read(content);
                            DataBufferUtils.release(dataBuffer);

                            String responseResult = readResponseBody(exchange, content);
                            responseResult = responseResult.substring(0, Math.min(responseResult.length(), MAX_RESPONSE_BODY_LOG_LENGTH));
                            loggingResponse.setBody(responseResult);
                            return bufferFactory.wrap(content);
                        }));
                    }
                }
                return super.writeWith(body);
            }
        };
    }

    private boolean isSuccessResponse(ServerHttpResponse response) {
        HttpStatusCode statusCode = response.getStatusCode();
        return statusCode == null || statusCode.is2xxSuccessful();
    }

    private boolean isJsonResponse(ServerWebExchange exchange, ServerHttpResponse response) {
        String originalResponseContentType = exchange.getAttribute(ServerWebExchangeUtils.ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR);
        if (StringUtils.isNotBlank(originalResponseContentType) && originalResponseContentType.contains(MediaType.APPLICATION_JSON_VALUE)) {
            return true;
        }
        MediaType contentType = response.getHeaders().getContentType();
        return contentType != null && MediaType.APPLICATION_JSON.isCompatibleWith(contentType);
    }

    private String readResponseBody(ServerWebExchange exchange, byte[] content) {
        if (!isGzipResponse(exchange, content)) {
            return new String(content, StandardCharsets.UTF_8);
        }
        try (ByteArrayOutputStream out = new ByteArrayOutputStream();
             GZIPInputStream in = new GZIPInputStream(new ByteArrayInputStream(content))) {
            byte[] buffer = new byte[1024];
            int len;
            while ((len = in.read(buffer)) > 0) {
                out.write(buffer, 0, len);
            }
            return out.toString(StandardCharsets.UTF_8);
        } catch (IOException e) {
            log.warn("Failed to read gzip response body for logging: {}", e.getMessage());
            return "";
        }
    }

    private boolean isGzipResponse(ServerWebExchange exchange, byte[] content) {
        boolean gzipHeader = exchange.getResponse().getHeaders().getOrEmpty(HttpHeaders.CONTENT_ENCODING)
                .stream()
                .anyMatch(value -> StringUtils.isNotBlank(value) && value.toLowerCase().contains("gzip"));
        return gzipHeader || hasGzipMagic(content);
    }

    private boolean hasGzipMagic(byte[] content) {
        return content != null
                && content.length >= 2
                && (content[0] & 0xff) == 0x1f
                && (content[1] & 0xff) == 0x8b;
    }

    /**
     * 请求装饰器，重新计算 headers.
     */
    @SuppressWarnings("all")
    private ServerHttpRequestDecorator requestDecorate(ServerWebExchange exchange,
                                                       HttpHeaders headers,
                                                       CachedBodyOutputMessage outputMessage) {
        return new ServerHttpRequestDecorator(exchange.getRequest()) {
            @Override
            public HttpHeaders getHeaders() {
                long contentLength = headers.getContentLength();
                HttpHeaders httpHeaders = new HttpHeaders();
                httpHeaders.putAll(super.getHeaders());
                setRequestTraceId(httpHeaders);
                if (contentLength > 0) {
                    httpHeaders.setContentLength(contentLength);
                } else {
                    httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
                }
                return httpHeaders;
            }

            @Override
            public Flux<DataBuffer> getBody() {
                return outputMessage.getBody();
            }
        };
    }

    /**
     * 必须再创建http req前设置, 否则无法修改.
     */
    private void setRequestTraceId(HttpHeaders httpHeaders) {
        try {
            httpHeaders.set(ServiceTraceId.REQ_INNER_TRACE_ID, ServiceTraceId.genTraceId());
        } catch (Exception ex) {
            log.error("TraceIdSetError: {}", Throwables.getStackTraceAsString(ex));
        }
    }


    /**
     * 顺序必须是<-1，否则标准的NettyWriteResponseFilter将在您的过滤器得到一个被调用的机会之前发送响应 也就是说如果不小于 -1 ，将不会执行获取后端响应的逻辑.
     */
    @Override
    public int getOrder() {
        return -100;
    }

}
