package io.scalecube.gateway.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.util.ByteBufferBackedInputStream;

import reactor.core.publisher.Flux;
import reactor.core.publisher.ReplayProcessor;

import io.netty.buffer.ByteBufAllocator;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.buffer.NettyDataBuffer;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.client.ReactorNettyWebSocketClient;
import org.springframework.web.util.UriComponentsBuilder;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.URI;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.List;

public class TestWebsocketClient {
  private static final Logger LOGGER = LoggerFactory.getLogger(TestWebsocketClient.class);
  private final InetSocketAddress serverAddress;
  private final URI uri;
  private final ReactorNettyWebSocketClient client;
  private final NettyDataBufferFactory BUFFER_FACTORY =
      new NettyDataBufferFactory(ByteBufAllocator.DEFAULT);

  public TestWebsocketClient(InetSocketAddress serverAddress, String path) {
    this.serverAddress = serverAddress;
    this.uri = getUri(path);
    client = new ReactorNettyWebSocketClient();
  }

  public <T> Flux<T> send(List<T> requests, Class<T> clazz) {
    ReplayProcessor<T> responses = ReplayProcessor.create();
    client.execute(uri,
        session -> {
          LOGGER.info("Start sending messages");

          return session
              .send(Flux.fromIterable(requests).map(this::encode))
              .thenMany(session.receive()
                  .map(nxt -> this.decode(nxt, clazz))
                  .take(requests.size())
                  .subscribeWith(responses))
              .then();

        })
        .doOnSuccessOrError((aVoid, ex) -> LOGGER.debug("Done: " + (ex != null ? ex.getMessage() : "success")))
        .block(Duration.ofSeconds(10));
    return responses;
  }

  private URI getUri(String path) {
    UriComponentsBuilder builder = UriComponentsBuilder
        .newInstance()
        .scheme("ws")
        .host(serverAddress.getAddress().getHostAddress())
        .port(serverAddress.getPort());

    if (path != null) {
      builder.path(path);
    }

    return builder.build().toUri();
  }

  private WebSocketMessage encode(Object message) {
    try {
      ByteArrayOutputStream baos = new ByteArrayOutputStream();
      new ObjectMapper().writeValue(baos, message);
      NettyDataBuffer dataBuffer = BUFFER_FACTORY.allocateBuffer();
      dataBuffer.write(baos.toByteArray());
      return new WebSocketMessage(WebSocketMessage.Type.BINARY, dataBuffer);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  }

  private <T> T decode(WebSocketMessage message, Class<T> clazz) {
    ByteBuffer buffer = message.getPayload().asByteBuffer();
    try {
      return new ObjectMapper().readValue(new ByteBufferBackedInputStream(buffer), clazz);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  }
}

@@ -195,16 +195,52 @@
 r.Flux;%0A
+import reactor.core.publisher.Mono;%0A
 import r
@@ -274,24 +274,24 @@
 yProcessor;%0A
-
 %0Aimport io.n
@@ -317,24 +317,62 @@
 Allocator;%0A%0A
+import org.reactivestreams.Publisher;%0A
 import org.s
@@ -921,24 +921,54 @@
 e.Duration;%0A
+import java.util.Collections;%0A
 import java.
@@ -1547,16 +1547,512 @@
 );%0A  %7D%0A%0A
+  public %3CT%3E Flux send(Publisher%3CT%3E requests, Class%3CT%3E clazz) %7B%0A    if (requests instanceof Flux) %7B%0A      return send(((Flux) requests).collectList(), clazz);%0A    %7D else if (requests instanceof Mono) %7B%0A      return send(Collections.singletonList(((Mono%3CT%3E) requests).block()), clazz);%0A    %7D%0A    throw new IllegalArgumentException(%22The request must be either Flux or Mono%22);%0A  %7D%0A%0A  public %3CT%3E Flux send(T request, Class%3CT%3E clazz) %7B%0A    return send(Collections.singletonList(request), clazz);%0A  %7D%0A%0A
   public
@@ -2211,16 +2211,16 @@
 on -%3E %7B%0A
+
         
@@ -2260,17 +2260,16 @@
 ages%22);%0A
-%0A
         

fix CR comments, fixes according to the fresh services snapshot