/*
 * JBoss, Home of Professional Open Source.
 * Copyright 2014 Red Hat, Inc., and individual contributors
 * as indicated by the @author tags.
 *
 * 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 io.undertow.websockets.core;

import org.xnio.ChannelListener;
import org.xnio.IoUtils;
import org.xnio.Pooled;

import java.io.IOException;
import java.nio.ByteBuffer;

/**
 * A receive listener that performs a callback when it receives a message
 *
 * @author Stuart Douglas
 */
public abstract class AbstractReceiveListener implements ChannelListener<WebSocketChannel> {

    @Override
    public void handleEvent(WebSocketChannel channel) {
        try {
            final StreamSourceFrameChannel result = channel.receive();
            if (result == null) {
                return;
            } else if (result.getType() == WebSocketFrameType.BINARY) {
                onBinary(channel, result);
            } else if (result.getType() == WebSocketFrameType.TEXT) {
                onText(channel, result);
            } else if (result.getType() == WebSocketFrameType.PONG) {
                onPong(channel, result);
            } else if (result.getType() == WebSocketFrameType.PING) {
                onPing(channel, result);
            } else if (result.getType() == WebSocketFrameType.CLOSE) {
                onClose(channel, result);
            }
        } catch (IOException e) {
            onError(channel, e);
        }
    }

    protected void onPing(WebSocketChannel webSocketChannel, StreamSourceFrameChannel channel) throws IOException {
        bufferFullMessage(channel);
    }

    protected void onClose(WebSocketChannel webSocketChannel, StreamSourceFrameChannel channel) throws IOException {
        bufferFullMessage(channel);
    }

    protected void onPong(WebSocketChannel webSocketChannel, StreamSourceFrameChannel messageChannel) throws IOException {
        bufferFullMessage(messageChannel);
    }

    protected void onText(WebSocketChannel webSocketChannel, StreamSourceFrameChannel messageChannel) throws IOException {
        bufferFullMessage(messageChannel);
    }

    protected void onBinary(WebSocketChannel webSocketChannel, StreamSourceFrameChannel messageChannel) throws IOException {
        bufferFullMessage(messageChannel);
    }

    protected void onError(WebSocketChannel channel, Throwable error) {
        IoUtils.safeClose(channel);
    }

    /**
     * Utility method that reads a full text or binary message, including all fragmented parts. Once the full message is
     * read then the {@link #onFullTextMessage(WebSocketChannel, BufferedTextMessage)} or
     * {@link #onFullBinaryMessage(WebSocketChannel, BufferedBinaryMessage)} method will be invoked.
     *
     * @param messageChannel The message channel
     */
    protected final void bufferFullMessage(StreamSourceFrameChannel messageChannel) {
        if (messageChannel.getType() == WebSocketFrameType.TEXT) {
            readBufferedText(messageChannel, new BufferedTextMessage(getMaxTextBufferSize(), true));
        } else if (messageChannel.getType() == WebSocketFrameType.BINARY) {
            readBufferedBinary(messageChannel, false, new BufferedBinaryMessage(getMaxBinaryBufferSize(), true));
        } else if (messageChannel.getType() == WebSocketFrameType.PONG) {
            readBufferedBinary(messageChannel, true, new BufferedBinaryMessage(getMaxPongBufferSize(), true));
        } else if (messageChannel.getType() == WebSocketFrameType.PING) {
            readBufferedBinary(messageChannel, true, new BufferedBinaryMessage(getMaxPingBufferSize(), true));
        } else if (messageChannel.getType() == WebSocketFrameType.CLOSE) {
            readBufferedBinary(messageChannel, true, new BufferedBinaryMessage(getMaxCloseBufferSize(), true));
        }
    }

    protected long getMaxBinaryBufferSize() {
        return -1;
    }

    protected long getMaxPongBufferSize() {
        return -1;
    }

    protected long getMaxCloseBufferSize() {
        return -1;
    }

    protected long getMaxPingBufferSize() {
        return -1;
    }

    protected long getMaxTextBufferSize() {
        return -1;
    }

    private void readBufferedBinary(final StreamSourceFrameChannel messageChannel, final boolean controlFrame, final BufferedBinaryMessage buffer) {

        buffer.read(messageChannel, new WebSocketCallback<BufferedBinaryMessage>() {
            @Override
            public void complete(WebSocketChannel channel, BufferedBinaryMessage context) {
                try {
                    WebSocketFrameType type = messageChannel.getType();
                    if (!controlFrame) {
                        onFullBinaryMessage(channel, buffer);
                    } else if (type == WebSocketFrameType.PONG) {
                        onFullPongMessage(channel, buffer);
                    } else if (type == WebSocketFrameType.PING) {
                        onFullPingMessage(channel, buffer);
                    } else if (type == WebSocketFrameType.CLOSE) {
                        onFullCloseMessage(channel, buffer);
                    }
                } catch (IOException e) {
                    AbstractReceiveListener.this.onError(channel, e);
                }
            }

            @Override
            public void onError(WebSocketChannel channel, BufferedBinaryMessage context, Throwable throwable) {
                context.getData().free();
                AbstractReceiveListener.this.onError(channel, throwable);
            }
        });
    }

    private void readBufferedText(StreamSourceFrameChannel messageChannel, final BufferedTextMessage textMessage) {
        textMessage.read(messageChannel, new WebSocketCallback<BufferedTextMessage>() {
            @Override
            public void complete(WebSocketChannel channel, BufferedTextMessage context) {
                try {
                    onFullTextMessage(channel, textMessage);
                } catch (IOException e) {
                    AbstractReceiveListener.this.onError(channel, e);
                }
            }

            @Override
            public void onError(WebSocketChannel channel, BufferedTextMessage context, Throwable throwable) {
                AbstractReceiveListener.this.onError(channel, throwable);
            }
        });
    }

    protected void onFullTextMessage(final WebSocketChannel channel, BufferedTextMessage message) throws IOException {

    }

    protected void onFullBinaryMessage(final WebSocketChannel channel, BufferedBinaryMessage message) throws IOException {

    }

    protected void onFullPingMessage(final WebSocketChannel channel, BufferedBinaryMessage message) throws IOException {
        final Pooled<ByteBuffer[]> data = message.getData();
        WebSockets.sendPong(data.getResource(), channel, new FreeDataCallback(data));
    }

    protected void onFullPongMessage(final WebSocketChannel channel, BufferedBinaryMessage message) throws IOException {
    }

    protected void onFullCloseMessage(final WebSocketChannel channel, BufferedBinaryMessage message) throws IOException {
        Pooled<ByteBuffer[]> data = message.getData();
        try {
            CloseMessage cm = new CloseMessage(data.getResource());
            onCloseMessage(cm, channel);
            if (!channel.isCloseFrameSent()) {
                WebSockets.sendClose(cm, channel, null);
            }
        } finally {
            data.free();
        }
    }

    protected void onCloseMessage(CloseMessage cm, WebSocketChannel channel) {

    }

    private static class FreeDataCallback implements WebSocketCallback<Void> {
        private final Pooled<ByteBuffer[]> data;

        public FreeDataCallback(Pooled<ByteBuffer[]> data) {
            this.data = data;
        }

        @Override
        public void complete(WebSocketChannel channel, Void context) {
            data.free();
        }

        @Override
        public void onError(WebSocketChannel channel, Void context, Throwable throwable) {
            data.free();
        }
    }
}

@@ -6941,33 +6941,32 @@
 s IOException %7B%0A
-%0A
     %7D%0A%0A    prote
@@ -7079,17 +7079,50 @@
 ption %7B%0A
+        message.getData().free();
 %0A
-
     %7D%0A%0A 
@@ -7508,32 +7508,66 @@
 s IOException %7B%0A
+        message.getData().free();%0A
     %7D%0A%0A    prote
@@ -8029,32 +8029,32 @@
        %7D%0A    %7D%0A%0A
+
     protected vo
@@ -8112,25 +8112,24 @@
  channel) %7B%0A
-%0A
     %7D%0A%0A    p

Make the default implementation free the data