package io.scalecube.services.transport.rsocket;

import io.rsocket.AbstractRSocket;
import io.rsocket.ConnectionSetupPayload;
import io.rsocket.Payload;
import io.rsocket.RSocket;
import io.rsocket.SocketAcceptor;
import io.rsocket.util.ByteBufPayload;
import io.scalecube.services.api.ServiceMessage;
import io.scalecube.services.exceptions.BadRequestException;
import io.scalecube.services.exceptions.ServiceException;
import io.scalecube.services.exceptions.ServiceUnavailableException;
import io.scalecube.services.methods.ServiceMethodInvoker;
import io.scalecube.services.methods.ServiceMethodRegistry;
import io.scalecube.services.transport.api.ReferenceCountUtil;
import io.scalecube.services.transport.api.ServiceMessageCodec;
import java.util.Optional;
import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * RSocket service acceptor. Implementation of {@link SocketAcceptor}. See for details and supported
 * methods -- {@link AbstractRSocket0}.
 */
public class RSocketServiceAcceptor implements SocketAcceptor {

  private static final Logger LOGGER = LoggerFactory.getLogger(RSocketServiceAcceptor.class);

  private final ServiceMessageCodec messageCodec;
  private final ServiceMethodRegistry methodRegistry;

  public RSocketServiceAcceptor(ServiceMessageCodec codec, ServiceMethodRegistry methodRegistry) {
    this.messageCodec = codec;
    this.methodRegistry = methodRegistry;
  }

  @Override
  public Mono<RSocket> accept(ConnectionSetupPayload setup, RSocket socket) {
    LOGGER.info("Accepted rSocket: {}, connectionSetup: {}", socket, setup);
    return Mono.just(new AbstractRSocket0());
  }

  private class AbstractRSocket0 extends AbstractRSocket {

    @Override
    public Mono<Payload> requestResponse(Payload payload) {
      return Mono.fromCallable(() -> toMessage(payload))
          .doOnNext(this::validateRequest)
          .flatMap(
              message -> {
                ServiceMethodInvoker methodInvoker = methodRegistry.getInvoker(message.qualifier());
                return methodInvoker.invokeOne(message, ServiceMessageCodec::decodeData);
              })
          .map(this::toPayload);
    }

    @Override
    public Flux<Payload> requestStream(Payload payload) {
      return Mono.fromCallable(() -> toMessage(payload))
          .doOnNext(this::validateRequest)
          .flatMapMany(
              message -> {
                ServiceMethodInvoker methodInvoker = methodRegistry.getInvoker(message.qualifier());
                return methodInvoker.invokeMany(message, ServiceMessageCodec::decodeData);
              })
          .map(this::toPayload);
    }

    @Override
    public Flux<Payload> requestChannel(Publisher<Payload> payloads) {
      return Flux
          .from(payloads)
          .map(this::toMessage)
          .switchOnFirst((first, messages) -> {
              if (first.hasValue()) {
                  ServiceMessage message = first.get();
                  validateRequest(message);
                  ServiceMethodInvoker methodInvoker = methodRegistry.getInvoker(message.qualifier());
                  return methodInvoker.invokeBidirectional(messages, ServiceMessageCodec::decodeData);
              }

              return messages;
          })
          .map(this::toPayload);
    }

    private Payload toPayload(ServiceMessage response) {
      return messageCodec.encodeAndTransform(response, ByteBufPayload::create);
    }

    private ServiceMessage toMessage(Payload payload) {
      return messageCodec.decode(payload.sliceData(), payload.sliceMetadata());
    }

    /**
     * Performs basic validation on incoming message: qualifier must be present, method invoker must
     * be present in the method registry by incoming qualifier. May throw exception
     *
     * @param message incoming message
     * @throws ServiceException in case qualfier is missing or method invoker is missing by given
     *     qualifier
     */
    private void validateRequest(ServiceMessage message) throws ServiceException {
      if (message.qualifier() == null) {
        Optional.ofNullable(message.data())
            .ifPresent(ReferenceCountUtil::safestRelease); // release message data if any
        LOGGER.error("Failed to invoke service with msg={}: qualifier is null", message);
        throw new BadRequestException("Qualifier is null in service msg request: " + message);
      }

      if (!methodRegistry.containsInvoker(message.qualifier())) {
        Optional.ofNullable(message.data())
            .ifPresent(ReferenceCountUtil::safestRelease); // release message data if any
        LOGGER.error(
            "Failed to invoke service with msg={}: no service invoker found by qualifier={}",
            message,
            message.qualifier());
        throw new ServiceUnavailableException(
            "No service invoker found by qualifier=" + message.qualifier());
      }
    }
  }
}

@@ -43,16 +43,138 @@
 ocket;%0A%0A
+import java.util.Optional;%0Aimport org.reactivestreams.Publisher;%0Aimport org.slf4j.Logger;%0Aimport org.slf4j.LoggerFactory;%0A
 import i
@@ -856,130 +856,8 @@
 ec;%0A
-import java.util.Optional;%0Aimport org.reactivestreams.Publisher;%0Aimport org.slf4j.Logger;%0Aimport org.slf4j.LoggerFactory;%0A
 impo
@@ -2837,27 +2837,16 @@
 urn Flux
-%0A          
 .from(pa
@@ -2910,16 +2910,31 @@
 OnFirst(
+%0A              
 (first, 
@@ -2944,24 +2944,26 @@
 sages) -%3E %7B%0A
+  
             
@@ -3134,32 +3134,54 @@
  methodInvoker =
+%0A                     
  methodRegistry.
@@ -3272,16 +3272,39 @@
 ctional(
+%0A                      
 messages
@@ -3349,27 +3349,31 @@
             
+  
 %7D%0A%0A
+  
             
@@ -3391,16 +3391,20 @@
 ssages;%0A
+    
         

format code