package io.scalecube.gateway.clientsdk.websocket;

import io.netty.buffer.ByteBuf;
import io.scalecube.gateway.clientsdk.ClientMessage;
import io.scalecube.gateway.clientsdk.ClientSettings;
import io.scalecube.gateway.clientsdk.ClientTransport;
import io.scalecube.gateway.clientsdk.codec.ClientMessageCodec;
import java.net.InetSocketAddress;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.MonoSink;
import reactor.core.scheduler.Schedulers;
import reactor.ipc.netty.http.client.HttpClient;
import reactor.ipc.netty.resources.LoopResources;

public final class WebsocketClientTransport implements ClientTransport {

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

  private static final AtomicReferenceFieldUpdater<WebsocketClientTransport, Mono>
      websocketMonoUpdater =
          AtomicReferenceFieldUpdater.newUpdater(
              WebsocketClientTransport.class, Mono.class, "websocketMono");

  private final ClientMessageCodec<ByteBuf> messageCodec;
  private final InetSocketAddress address;
  private final HttpClient httpClient;
  private final AtomicLong sidCounter = new AtomicLong();

  private volatile Mono<?> websocketMono;

  /**
   * Creates instance of websocket client transport.
   *
   * @param settings client settings
   * @param messageCodec client message codec
   * @param loopResources loop resources
   */
  public WebsocketClientTransport(
      ClientSettings settings,
      ClientMessageCodec<ByteBuf> messageCodec,
      LoopResources loopResources) {
    this.messageCodec = messageCodec;

    address = InetSocketAddress.createUnresolved(settings.host(), settings.port());

    httpClient =
        HttpClient.create(
            options ->
                options
                    .disablePool()
                    .compression(false)
                    .connectAddress(() -> address)
                    .loopResources(loopResources));
  }

  @Override
  public Mono<ClientMessage> requestResponse(ClientMessage request) {
    return Mono.defer(
        () ->
            getOrConnect()
                .flatMap(
                    session ->
                        Mono.create(
                            sink -> {
                              String sid = String.valueOf(sidCounter.incrementAndGet());

                              sink.onCancel(() -> session.send(cancelMessage(sid)).subscribe());

                              session
                                  .send(enrichForSend(requestMessage(request, sid)))
                                  .then(session.receive(sid).singleOrEmpty())
                                  .publishOn(Schedulers.parallel())
                                  .map(this::enrichForRecv)
                                  .subscribe(sink::success, sink::error, sink::success);
                            })));
  }

  @Override
  public Flux<ClientMessage> requestStream(ClientMessage request) {
    return Flux.defer(
        () ->
            getOrConnect()
                .flatMapMany(
                    session ->
                        Flux.create(
                            sink -> {
                              String sid = String.valueOf(sidCounter.incrementAndGet());

                              sink.onCancel(() -> session.send(cancelMessage(sid)).subscribe());

                              session
                                  .send(enrichForSend(requestMessage(request, sid)))
                                  .thenMany(session.receive(sid))
                                  .publishOn(Schedulers.parallel())
                                  .map(this::enrichForRecv)
                                  .subscribe(sink::next, sink::error, sink::complete);
                            })));
  }

  @Override
  public Mono<Void> close() {
    return Mono.defer(
        () -> {
          // noinspection unchecked
          Mono<WebsocketSession> curr = websocketMonoUpdater.get(this);
          return curr == null ? Mono.empty() : curr.flatMap(WebsocketSession::close);
        });
  }

  private Mono<WebsocketSession> getOrConnect() {
    // noinspection unchecked
    return Mono.defer(() -> websocketMonoUpdater.updateAndGet(this, this::getOrConnect0));
  }

  private Mono<WebsocketSession> getOrConnect0(Mono<WebsocketSession> prev) {
    if (prev != null) {
      return prev;
    }

    return httpClient
        .ws("/")
        .flatMap(
            response ->
                Mono.create(
                    (MonoSink<WebsocketSession> sink) ->
                        response
                            .receiveWebsocket(
                                (in, out) -> {
                                  LOGGER.info("Connected successfully to {}", address);

                                  WebsocketSession session =
                                      new WebsocketSession(in, out, messageCodec);

                                  sink.success(session);

                                  return session.onClose(
                                      () ->
                                          LOGGER.info(
                                              "Connection to {} has been closed successfully",
                                              address));
                                })
                            .doOnError(sink::error)
                            .subscribe()))
        .doOnError(
            throwable -> {
              LOGGER.warn("Connection to {} is failed, cause: {}", address, throwable);
              websocketMonoUpdater.getAndSet(this, null);
            })
        .cache();
  }

  private ClientMessage enrichForSend(ClientMessage clientMessage) {
    return ClientMessage.from(clientMessage)
        .header("client-send-time", String.valueOf(System.currentTimeMillis()))
        .build();
  }

  private ClientMessage enrichForRecv(ClientMessage message) {
    return ClientMessage.from(message)
        .header("client-recv-time", String.valueOf(System.currentTimeMillis()))
        .build();
  }

  private ClientMessage requestMessage(ClientMessage request, String sid) {
    return ClientMessage.from(request).header("sid", sid).build();
  }

  private ClientMessage cancelMessage(String sid) {
    return ClientMessage.builder()
        .header("sid", sid)
        .header("sig", Signal.CANCEL.codeAsString())
        .build();
  }
}

@@ -509,40 +509,8 @@
 ry;%0A
-import reactor.core.Disposable;%0A
 impo
@@ -2360,90 +2360,11 @@
 n -%3E
-%0A                        Mono.create(%0A                            sink -%3E %7B%0A      
+ %7B%0A
     
@@ -2373,34 +2373,32 @@
                 
-  
 String sid = Str
@@ -2432,33 +2432,32 @@
 ementAndGet());%0A
-%0A
                 
@@ -2466,113 +2466,14 @@
     
-        sink.onCancel(() -%3E session.send(cancelMessage(sid)).subscribe());%0A%0A                             
+return
  ses
@@ -2473,32 +2473,24 @@
 urn session%0A
-        
             
@@ -2576,32 +2576,24 @@
             
-  
-      
 .then(sessio
@@ -2646,32 +2646,24 @@
             
-        
 .publishOn(S
@@ -2706,32 +2706,24 @@
             
-      
-  
 .map(this::e
@@ -2766,68 +2766,70 @@
     
-        .subscribe(sink::success, sink::error, sink::success
+.doOnCancel(() -%3E session.send(cancelMessage(sid)).subscribe()
 );%0A 
@@ -2843,34 +2843,25 @@
             
-        %7D)
+%7D
 ));%0A  %7D%0A%0A  @
@@ -3065,92 +3065,11 @@
 n -%3E
-%0A                        Flux.create(%0A                            sink -%3E %7B%0A        
+ %7B%0A
     
@@ -3141,25 +3141,24 @@
 tAndGet());%0A
-%0A
             
@@ -3171,113 +3171,14 @@
     
-        sink.onCancel(() -%3E session.send(cancelMessage(sid)).subscribe());%0A%0A                             
+return
  ses
@@ -3178,32 +3178,24 @@
 urn session%0A
-        
             
@@ -3281,32 +3281,24 @@
             
-      
-  
 .thenMany(se
@@ -3317,24 +3317,16 @@
 e(sid))%0A
-        
         
@@ -3399,32 +3399,24 @@
             
-        
 .map(this::e
@@ -3459,66 +3459,70 @@
     
-        .subscribe(sink::next, sink::error, sink::complete
+.doOnCancel(() -%3E session.send(cancelMessage(sid)).subscribe()
 );%0A 
@@ -3540,26 +3540,17 @@
         
-        %7D)
+%7D
 ));%0A  %7D%0A

Remove redundant Mono.create and Flux.create