/**
 * Copyright 2008-2009. Chongqing Communications Industry Services Co.,Ltd Information Technology Branch. All rights reserved. <a>http://www.crunii.com</a>
 */
package com.crunii.micro.gateway.filter.web.dt;

import com.crunii.micro.common.utils.BusiLogUtil;
import org.apache.commons.io.IOUtils;
import org.reactivestreams.Publisher;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Map;

/**
 * @author 田平 create 2020年1月15日下午4:41:41
 */
public class ResponseLoggingDecorator extends ServerHttpResponseDecorator {

	public static final String RSP_BODY_KEY = ResponseLoggingDecorator.class.getName() + ".rsp.body.key";

	private Map<String, Object> attributes;

	public ResponseLoggingDecorator(ServerHttpResponse delegate, Map<String, Object> attributes) {
		super(delegate);
		this.attributes = attributes;
	}

	@Override
	public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
		if (BusiLogUtil.checkContentType(getHeaders().getContentType())) {
			ByteArrayOutputStream payload = new ByteArrayOutputStream(1024);
			return super.writeWith(Flux.from(body)
					.doOnNext(buffer -> {
						byte[] bytes = new byte[buffer.readableByteCount()];
						int position = buffer.readPosition();
						buffer.read(bytes);
						buffer.readPosition(position);
						try {
							payload.write(bytes);
						} catch (IOException e) {
							IOUtils.closeQuietly(payload);
							throw new RuntimeException("LoggingJsonDecoder.decodeToMono error!", e);
						}
					})).doOnSuccess(v -> {
				String rspBody = new String(payload.toByteArray(), StandardCharsets.UTF_8);
				attributes.put(RSP_BODY_KEY, rspBody);
				IOUtils.closeQuietly(payload);
			}).doOnError(e -> {
				String rspBody = new String(payload.toByteArray(), StandardCharsets.UTF_8);
				attributes.put(RSP_BODY_KEY, rspBody);
				IOUtils.closeQuietly(payload);
			});
		} else {
			return super.writeWith(body);
		}
	}
}
