package test;

import com.google.common.collect.Lists;
import com.google.common.hash.Hashing;
import com.google.common.util.concurrent.SettableFuture;
import io.undertow.Undertow;
import io.undertow.servlet.api.DeploymentManager;
import io.undertow.websockets.jsr.WebSocketDeploymentInfo;
import org.xnio.OptionMap;
import org.xnio.Xnio;
import org.xnio.XnioWorker;

import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.net.URI;
import java.nio.ByteBuffer;
import java.util.List;
import java.util.Random;

import static io.undertow.servlet.Servlets.defaultContainer;
import static io.undertow.servlet.Servlets.deployment;
import static io.undertow.websockets.jsr.WebSocketDeploymentInfo.ATTRIBUTE_NAME;

public class Test {

    private static int MESSAGES = 1000;


    @ServerEndpoint("/webSocket")
    public static class EchoSocket {
        private List<String> echos = Lists.newLinkedList();
        public static final SettableFuture<List<String>> receivedEchos = SettableFuture.create();

        @OnMessage
        public void onMessage(ByteBuffer dataBuffer, Session session) throws IOException {
            byte[] hd = new byte[dataBuffer.remaining()];
            dataBuffer.get(hd);
            String hash = Hashing.crc32().hashBytes(hd).toString();
            echos.add(hash);
            if (echos.size() == MESSAGES)
                receivedEchos.set(echos);
            session.getBasicRemote().sendBinary(dataBuffer);
        }

    }

    public static void main(String[] args) throws Exception {
        final Xnio xnio = Xnio.getInstance("nio", Undertow.class.getClassLoader());
        final XnioWorker xnioWorker = xnio.createWorker(OptionMap.builder().getMap());
        final WebSocketDeploymentInfo webSockets = new WebSocketDeploymentInfo()
                .addEndpoint(EchoSocket.class)
                .setWorker(xnioWorker);
        final DeploymentManager deployment = defaultContainer()
                .addDeployment(deployment()
                                       .setClassLoader(Test.class.getClassLoader())
                                       .setContextPath("/")
                                       .setDeploymentName("embedded-websockets")
                                       .addServletContextAttribute(ATTRIBUTE_NAME, webSockets)
                );

        deployment.deploy();


        final Undertow server = Undertow.builder()
                .addHttpListener(8080, "localhost")
                .setHandler(deployment.start())
                .build();
        server.start();
        final ClientEndpointConfig clientEndpointConfig = ClientEndpointConfig.Builder.create().build();
        ContainerProvider.getWebSocketContainer()
                .connectToServer(new Endpoint() {
                                     @Override
                                     public void onOpen(final Session session, EndpointConfig endpointConfig) {

                                         try {
                                             RemoteEndpoint.Basic rem = session.getBasicRemote();
                                             List<String> messages = Lists.newLinkedList();
                                             for (int i = 0; i < MESSAGES; i++) {
                                                 byte[] data = new byte[2048];
                                                 (new Random()).nextBytes(data);
                                                 String crc = Hashing.crc32().hashBytes(data).toString();
                                                 System.out.println((1 + i) + ": " + crc);
                                                 rem.sendBinary(ByteBuffer.wrap(data));
                                                 messages.add(crc);
                                             }
                                             System.out.println("Send.");

                                             List<String> received = EchoSocket.receivedEchos.get();

                                             System.out.println("Match: " + messages.equals(received));
                                             for (int i = 0; i < messages.size(); i++) {
                                                 if (received.size() <= i) {
                                                     System.out.println(i + ": should be " + messages.get(i) + " but is empty.");
                                                 } else {
                                                     if (!messages.get(i).equals(received.get(i))) {
                                                         System.out.println(i + ": should be " + messages.get(i) + " but is " + received.get(i) + " (but found at " + received.indexOf(messages.get(i)) + ").");
                                                     }
                                                 }
                                             }
                                         } catch (Throwable t) {
                                             System.out.println(t);
                                         } finally {
                                             server.stop();
                                         }
                                     }


                                 }, clientEndpointConfig, new URI("ws://localhost:8080/webSocket")
                );
    }
}