/*
 * Copyright (c) 2011-2016 Pivotal Software Inc, All Rights Reserved.
 *
 * 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 reactor.io.netty.http;

import io.netty.buffer.ByteBuf;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelPromise;
import io.netty.handler.codec.http.FullHttpResponse;
import io.netty.handler.codec.http.HttpObjectAggregator;
import io.netty.handler.codec.http.HttpResponse;
import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
import io.netty.handler.codec.http.websocketx.CloseWebSocketFrame;
import io.netty.handler.codec.http.websocketx.PingWebSocketFrame;
import io.netty.handler.codec.http.websocketx.PongWebSocketFrame;
import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
import io.netty.handler.codec.http.websocketx.WebSocketClientHandshaker;
import io.netty.handler.codec.http.websocketx.WebSocketClientHandshakerFactory;
import io.netty.handler.codec.http.websocketx.WebSocketVersion;
import io.netty.util.ReferenceCountUtil;
import reactor.core.publisher.Flux;

/**
 * @author Stephane Maldini
 */
final class NettyWebSocketClientHandler extends NettyHttpClientHandler {

	final WebSocketClientHandshaker handshaker;

	final ChannelPromise handshakerResult;

	final boolean plainText;

	NettyWebSocketClientHandler(
			String protocols,
			NettyHttpClientHandler originalHandler,
			boolean plainText) {
		super(originalHandler.getHandler(), originalHandler.tcpStream, originalHandler.currentURI);
		this.httpChannel = originalHandler.httpChannel;
		this.replySubscriber = originalHandler.replySubscriber;
		this.plainText = plainText;

		handshaker =
				WebSocketClientHandshakerFactory.newHandshaker(currentURI,
						WebSocketVersion.V13,
						protocols,
						false,
						httpChannel.removeTransferEncodingChunked().headers());

		handshakerResult = httpChannel.delegate().newPromise();
		handshaker.handshake(httpChannel.delegate()).addListener(f -> {
			if(!f.isSuccess()) {
				handshakerResult.tryFailure(f.cause());
				return;
			}
			httpChannel.delegate().read();
		});
	}

	@Override
	protected void writeLast(ChannelHandlerContext ctx) {
		ctx.writeAndFlush(new CloseWebSocketFrame());
	}

	@Override
	protected void postRead(ChannelHandlerContext ctx, Object msg) {
		if(CloseWebSocketFrame.class.isAssignableFrom(msg.getClass())){
			ctx.channel().close();
		}
	}

	@Override
	@SuppressWarnings("unchecked")
	public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
		Class<?> messageClass = msg.getClass();
		if (FullHttpResponse.class.isAssignableFrom(messageClass)) {
			ctx.pipeline().remove(HttpObjectAggregator.class);
			HttpResponse response = (HttpResponse) msg;
			if (httpChannel != null) {
				httpChannel.setNettyResponse(response);
			}

			checkResponseCode(ctx, response);

			if(!handshaker.isHandshakeComplete()) {
				handshaker.finishHandshake(ctx.channel(), (FullHttpResponse) msg);
			}
			ctx.fireChannelRead(msg);
			handshakerResult.trySuccess();

			if(!discardBody && replySubscriber != null){
				Flux.just(httpChannel).subscribe(replySubscriber);
			}
			return;
		}
		if (PingWebSocketFrame.class.isAssignableFrom(messageClass)) {
			ctx.channel().writeAndFlush(new PongWebSocketFrame(((PingWebSocketFrame)msg)
					.content().retain()));
			return;
		}
		if (TextWebSocketFrame.class.isAssignableFrom(messageClass)) {
			try {
				if(channelSubscriber != null) {
					channelSubscriber.onNext(msg);
				}
			} finally {
				ReferenceCountUtil.release(msg);
			}
		} else if (CloseWebSocketFrame.class.isAssignableFrom(messageClass)) {
			ctx.close();
		} else {
			doRead(ctx, msg);
		}
	}

	@Override
	protected ChannelFuture doOnWrite(Object data, ChannelHandlerContext ctx) {
		if (data instanceof ByteBuf) {
			if (plainText){
				return ctx.write(new TextWebSocketFrame((ByteBuf)data));
			}
			return ctx.write(new BinaryWebSocketFrame((ByteBuf)data));
		}
		else if (data instanceof String) {
			return ctx.write(new TextWebSocketFrame((String)data));
		}
		else {
			return ctx.write(data);
		}
	}
}

@@ -2666,17 +2666,21 @@
 oid 
-writeLast
+doOnTerminate
 (Cha
@@ -2697,24 +2697,98 @@
 rContext ctx
+,%0A%09%09%09ChannelFuture last,%0A%09%09%09ChannelPromise promise,%0A%09%09%09Throwable exception
 ) %7B%0A%09%09ctx.wr
@@ -2786,32 +2786,24 @@
 %0A%09%09ctx.write
-AndFlush
 (new CloseWe
@@ -2819,16 +2819,70 @@
 ame());%0A
+%09%09super.doOnTerminate(ctx, last, promise, exception);%0A
 %09%7D%0A%0A%09@Ov

tweak body