/*
 * Copyright 2014-2015 Red Hat, Inc, and individual contributors.
 *
 * 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
 *
 * http://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.projectodd.wunderboss.web.async.websocket;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.HashSet;
import java.util.Set;

import io.undertow.server.HttpHandler;
import io.undertow.server.HttpServerExchange;
import io.undertow.server.handlers.ResponseCodeHandler;
import io.undertow.util.HeaderValues;
import io.undertow.util.Headers;
import io.undertow.websockets.WebSocketConnectionCallback;
import io.undertow.websockets.WebSocketProtocolHandshakeHandler;
import io.undertow.websockets.core.WebSocketChannel;
import io.undertow.websockets.core.AbstractReceiveListener;
import io.undertow.websockets.core.BufferedTextMessage;
import io.undertow.websockets.core.BufferedBinaryMessage;
import io.undertow.websockets.core.CloseMessage;
import io.undertow.websockets.spi.WebSocketHttpExchange;
import org.projectodd.wunderboss.web.async.Channel;
import org.xnio.Buffers;
import org.xnio.Pooled;


public class UndertowWebsocket {

    public static HttpHandler createHandler(final ThreadLocal<HttpServerExchange> requestTL,
                                            final WebsocketInitHandler checker,
                                            final HttpHandler next) {
        WebSocketConnectionCallback callback = new WebSocketConnectionCallback() {
            public void onConnect (WebSocketHttpExchange exchange, WebSocketChannel channel) {
                final DelegatingUndertowEndpoint endpoint = new DelegatingUndertowEndpoint();
                if (checker.shouldConnect(exchange, endpoint)) {
                    endpoint.onOpen(channel, exchange);
                    channel.getReceiveSetter().set(new AbstractReceiveListener() {
                        protected void onError (WebSocketChannel channel, Throwable error) {
                            endpoint.onError(channel, error);
                        }
                        protected void onCloseMessage (CloseMessage message, WebSocketChannel channel) {
                            endpoint.onClose(channel, message);
                        }
                        protected void onFullTextMessage (WebSocketChannel channel, BufferedTextMessage message) {
                            endpoint.onMessage(channel, message.getData());
                        }
                        protected void onFullBinaryMessage (WebSocketChannel channel, BufferedBinaryMessage message) {
                            Pooled<ByteBuffer[]> pooled = message.getData();
                            try {
                                ByteBuffer[] payload = pooled.getResource();
                                endpoint.onMessage(channel, toArray(payload));
                            } finally {
                                pooled.free();
                            }
                        }
                    });
                    channel.resumeReceives();
                } else {
                    try {
                        channel.close();
                    } catch (IOException e) {
                        throw new RuntimeException("Error closing websocket", e);
                    }
                }
            }
        };

        final HttpHandler downstream = next==null ? ResponseCodeHandler.HANDLE_404 : next;
        final HttpHandler wsHandler = new WebSocketProtocolHandshakeHandler(callback, downstream);


        return new HttpHandler() {
            @Override
            public void handleRequest(HttpServerExchange exchange) throws Exception {
                HeaderValues upgrade = exchange.getRequestHeaders().get(Headers.UPGRADE);
                if (upgrade != null && "websocket".equalsIgnoreCase(upgrade.peek())) {
                    requestTL.set(exchange);
                    try {
                        wsHandler.handleRequest(exchange);
                    } finally {
                        requestTL.remove();
                    }
                } else {
                    downstream.handleRequest(exchange);
                }
            }
        };


    }

    public static WebsocketChannel createWebsocketChannel(final Channel.OnOpen onOpen,
                                                          final Channel.OnError onError,
                                                          final Channel.OnClose onClose,
                                                          final WebsocketChannel.OnMessage onMessage) {
        Channel.OnClose removeOnClose = new Channel.OnClose() {
            @Override
            public void handle(final Channel channel, final Object code, final String reason) {
                openChannels.remove(channel);
                if (onClose != null) {
                    onClose.handle(channel, code, reason);
                }
            }
        };

        WebsocketChannel chan = new UndertowWebsocketChannel(onOpen, onError, removeOnClose, onMessage);

        openChannels.add(chan);

        return chan;
    }

    // Lifted from Undertow's FrameHandler.java
    protected static byte[] toArray(ByteBuffer... payload) {
        if (payload.length == 1) {
            ByteBuffer buf = payload[0];
            if (buf.hasArray() && buf.arrayOffset() == 0 && buf.position() == 0) {
                return buf.array();
            }
        }
        int size = (int) Buffers.remaining(payload);
        byte[] data = new byte[size];
        for (ByteBuffer buf : payload) {
            buf.get(data);
        }
        return data;
    }

    private final static Set<Channel> openChannels = new HashSet<>();
}

@@ -681,113 +681,8 @@
 t;%0A%0A
-import java.io.IOException;%0Aimport java.nio.ByteBuffer;%0Aimport java.util.HashSet;%0Aimport java.util.Set;%0A%0A
 impo
@@ -1017,61 +1017,8 @@
 er;%0A
-import io.undertow.websockets.core.WebSocketChannel;%0A
 impo
@@ -1112,28 +1112,30 @@
 ore.Buffered
-Text
+Binary
 Message;%0Aimp
@@ -1166,38 +1166,36 @@
 ts.core.Buffered
-Binary
+Text
 Message;%0Aimport 
@@ -1270,11 +1270,12 @@
 ets.
-spi
+core
 .Web
@@ -1284,121 +1284,178 @@
 cket
-HttpExchange;%0Aimport org.projectodd.wunderboss.web.async.Channel;%0Aimport org.xnio.Buffers;%0Aimport org.xnio.Pooled
+Channel;%0Aimport io.undertow.websockets.spi.WebSocketHttpExchange;%0Aimport org.xnio.Buffers;%0Aimport org.xnio.Pooled;%0A%0Aimport java.io.IOException;%0Aimport java.nio.ByteBuffer
 ;%0A%0A%0A
@@ -4562,914 +4562,8 @@
  %7D%0A%0A
-    public static WebsocketChannel createWebsocketChannel(final Channel.OnOpen onOpen,%0A                                                          final Channel.OnError onError,%0A                                                          final Channel.OnClose onClose,%0A                                                          final WebsocketChannel.OnMessage onMessage) %7B%0A        Channel.OnClose removeOnClose = new Channel.OnClose() %7B%0A            @Override%0A            public void handle(final Channel channel, final Object code, final String reason) %7B%0A                openChannels.remove(channel);%0A                if (onClose != null) %7B%0A                    onClose.handle(channel, code, reason);%0A                %7D%0A            %7D%0A        %7D;%0A%0A        WebsocketChannel chan = new UndertowWebsocketChannel(onOpen, onError, removeOnClose, onMessage);%0A%0A        openChannels.add(chan);%0A%0A        return chan;%0A    %7D%0A%0A
     
@@ -5085,78 +5085,7 @@
    %7D
-%0A%0A    private final static Set%3CChannel%3E openChannels = new HashSet%3C%3E();
 %0A%7D%0A

Remove unused method that was accidentally commited with last commit.