/*
 * Copyright 2020-2020 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.springframework.cloud.function.rsocket;

import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.function.Function;

import io.rsocket.Payload;
import io.rsocket.RSocket;
import io.rsocket.core.RSocketConnector;
import io.rsocket.transport.netty.client.TcpClientTransport;
import io.rsocket.util.DefaultPayload;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.reactivestreams.Publisher;
import reactor.util.retry.Retry;

import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry.FunctionInvocationWrapper;
import org.springframework.messaging.Message;
import org.springframework.messaging.support.MessageBuilder;



/**
 *
 * @author Oleg Zhurakousky
 * @since 3.1
 *
 */
class RSocketForwardingFunction implements Function<Message<byte[]>, Publisher<Message<byte[]>>> {
	private static Log logger = LogFactory.getLog(RSocketForwardingFunction.class);

	private final RSocket rSocket;

	private final FunctionInvocationWrapper targetFunction;

	RSocketForwardingFunction(FunctionInvocationWrapper targetFunction, InetSocketAddress outputAddress) {
		this.targetFunction = targetFunction;
		this.rSocket = outputAddress == null ? null
				: RSocketConnector.connectWith(TcpClientTransport.create(outputAddress))
					.log()
					.retryWhen(Retry.backoff(5, Duration.ofSeconds(1)))
					.block();
	}

	@SuppressWarnings("unchecked")
	@Override
	public Publisher<Message<byte[]>> apply(Message<byte[]> input) {
		if (logger.isDebugEnabled()) {
			logger.debug("Executiing: " + this.targetFunction);
		}

		Object rawResult = this.targetFunction.apply(input);
		Publisher<Message<byte[]>> resultMessage = this.rSocket
				.requestStream(DefaultPayload.create(((Message<byte[]>) rawResult).getPayload()))
				.map(this::buildResultMessage);
		return resultMessage;
	}

	private Message<byte[]> buildResultMessage(Payload payload) {
		ByteBuffer payloadBuffer = payload.getData();
		byte[] payloadData = new byte[payloadBuffer.remaining()];
		payloadBuffer.get(payloadData);
		return MessageBuilder.withPayload(payloadData).build();
	}
}

@@ -1110,24 +1110,60 @@
 .Publisher;%0A
+import reactor.core.publisher.Mono;%0A
 import react
@@ -1401,17 +1401,16 @@
 lder;%0A%0A%0A
-%0A
 /**%0A *%0A 
@@ -1436,16 +1436,42 @@
 akousky%0A
+ * @author Artem Bilan%0A *%0A
  * @sinc
@@ -1582,16 +1582,17 @@
 %5B%5D%3E%3E%3E %7B%0A
+%0A
 %09private
@@ -1603,18 +1603,24 @@
 tic 
-Log logger
+final Log LOGGER
  = L
@@ -1690,23 +1690,33 @@
 nal 
+Mono%3C
 RSocket
+%3E
  r
-S
+s
 ocket
+Mono
 ;%0A%0A%09
@@ -1928,17 +1928,24 @@
 is.r
-S
+s
 ocket
- = 
+Mono =%0A%09%09%09
 outp
@@ -1961,17 +1961,21 @@
  == null
- 
+%0A%09%09%09%09
 ? null%0A%09
@@ -2001,142 +2001,129 @@
 or.c
-onnectWith(TcpClientTransport.create(outputAddress))%0A%09%09%09%09%09.log()%0A%09%09%09%09%09.retryWhen(Retry.backoff(5, Duration.ofSeconds(1)))%0A%09%09%09%09%09.block(
+reate()%0A%09%09%09%09%09%09.reconnect(Retry.backoff(5, Duration.ofSeconds(1)))%0A%09%09%09%09%09%09.connect(TcpClientTransport.create(outputAddress)
 );%0A%09
@@ -2240,22 +2240,22 @@
 %7B%0A%09%09if (
-logger
+LOGGER
 .isDebug
@@ -2274,14 +2274,14 @@
 %0A%09%09%09
-logger
+LOGGER
 .deb
@@ -2291,17 +2291,16 @@
 %22Executi
-i
 ng: %22 + 
@@ -2387,68 +2387,72 @@
 ;%0A%09%09
-Publisher%3CMessage%3Cbyte%5B%5D%3E%3E resultMessage = this.rS
+return this.rsocketMono%0A%09%09%09.flatMapMany((rs
 ocket
+) -%3E
 %0A%09%09%09%09
+rsocket
 .req
@@ -2528,18 +2528,18 @@
 load()))
+)
 %0A
-%09
 %09%09%09.map(
@@ -2569,32 +2569,8 @@
 e);%0A
-%09%09return resultMessage;%0A
 %09%7D%0A%0A
@@ -2835,10 +2835,11 @@
 d();%0A%09%7D%0A
+%0A
 %7D%0A

Improve RSocketForwardingFunction