/*
 * Copyright 2014 The Vibe Project
 *
 * 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.atmosphere.vibe.platform.server.netty4;

import io.netty.channel.Channel;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.channel.ChannelPipeline;
import io.netty.handler.codec.http.DefaultFullHttpRequest;
import io.netty.handler.codec.http.FullHttpRequest;
import io.netty.handler.codec.http.HttpContent;
import io.netty.handler.codec.http.HttpHeaders;
import io.netty.handler.codec.http.HttpMethod;
import io.netty.handler.codec.http.HttpRequest;
import io.netty.handler.codec.http.HttpResponse;
import io.netty.handler.codec.http.HttpServerCodec;
import io.netty.handler.codec.http.LastHttpContent;
import io.netty.handler.codec.http.websocketx.WebSocketFrame;
import io.netty.handler.codec.http.websocketx.WebSocketServerHandshaker;
import io.netty.handler.codec.http.websocketx.WebSocketServerHandshakerFactory;
import io.netty.handler.ssl.SslHandler;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.atmosphere.vibe.platform.Action;
import org.atmosphere.vibe.platform.Actions;
import org.atmosphere.vibe.platform.SimpleActions;
import org.atmosphere.vibe.platform.server.ServerHttpExchange;
import org.atmosphere.vibe.platform.server.ServerWebSocket;

/**
 * Processes {@link HttpRequest} and {@link HttpResponse} into
 * {@link NettyServerHttpExchange} and {@link NettyServerWebSocket}. It requires
 * {@link HttpServerCodec} in front of this codec.
 *
 * @author Donghwan Kim
 */
public class VibeServerCodec extends ChannelInboundHandlerAdapter {

    private Actions<ServerHttpExchange> httpActions = new SimpleActions<>();
    private Actions<ServerWebSocket> wsActions = new SimpleActions<>();
    private Map<Channel, NettyServerHttpExchange> httpMap = new ConcurrentHashMap<>();
    private Map<Channel, NettyServerWebSocket> wsMap = new ConcurrentHashMap<>();
    private Map<Channel, FullHttpRequest> wsReqMap = new ConcurrentHashMap<>();

    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) {
        if (msg instanceof HttpRequest) {
            HttpRequest req = (HttpRequest) msg;
            if (!accept(req)) {
                return;
            }
            if (req.getMethod() == HttpMethod.GET && req.headers().contains(HttpHeaders.Values.UPGRADE, HttpHeaders.Values.WEBSOCKET, true)) {
                // Because WebSocketServerHandshaker requires FullHttpRequest
                FullHttpRequest wsRequest = new DefaultFullHttpRequest(req.getProtocolVersion(), req.getMethod(), req.getUri());
                wsRequest.headers().set(req.headers());
                wsReqMap.put(ctx.channel(), wsRequest);
            } else {
                NettyServerHttpExchange http = new NettyServerHttpExchange(ctx, req);
                httpMap.put(ctx.channel(), http);
                httpActions.fire(http);
            }
        } else if (msg instanceof HttpContent) {
            FullHttpRequest wsReq = wsReqMap.get(ctx.channel());
            if (wsReq != null) {
                wsReq.content().writeBytes(((HttpContent) msg).content());
                if (msg instanceof LastHttpContent) {
                    wsReqMap.remove(ctx.channel());
                    WebSocketServerHandshakerFactory factory = new WebSocketServerHandshakerFactory(getWebSocketLocation(ctx.pipeline(), wsReq), null, true);
                    WebSocketServerHandshaker handshaker = factory.newHandshaker(wsReq);
                    if (handshaker == null) {
                        WebSocketServerHandshakerFactory.sendUnsupportedVersionResponse(ctx.channel());
                    } else {
                        handshaker.handshake(ctx.channel(), wsReq);
                        NettyServerWebSocket ws = new NettyServerWebSocket(ctx, wsReq, handshaker);
                        wsMap.put(ctx.channel(), ws);
                        wsActions.fire(ws);
                    }
                }
            } else {
                NettyServerHttpExchange http = httpMap.get(ctx.channel());
                if (http != null) {
                    http.handleChunk((HttpContent) msg);
                }
            }
        } else if (msg instanceof WebSocketFrame) {
            NettyServerWebSocket ws = wsMap.get(ctx.channel());
            if (ws != null) {
                ws.handleFrame((WebSocketFrame) msg);
            }
        }
    }
    
    /**
     * Whether to delegate this request to Vibe or not. By default, it accepts
     * every request.
     */
    protected boolean accept(HttpRequest req) {
        return true;
    }

    private String getWebSocketLocation(ChannelPipeline pipeline, HttpRequest req) {
        return (pipeline.get(SslHandler.class) == null ? "ws://" : "wss://") + HttpHeaders.getHost(req) + req.getUri();
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        NettyServerHttpExchange http = httpMap.get(ctx.channel());
        if (http != null) {
            http.handleError(cause);
        }
        NettyServerWebSocket ws = wsMap.get(ctx.channel());
        if (ws != null) {
            ws.handleError(cause);
        }
        ctx.close();
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) {
        NettyServerHttpExchange http = httpMap.remove(ctx.channel());
        if (http != null) {
            http.handleClose();
        }
        NettyServerWebSocket ws = wsMap.remove(ctx.channel());
        if (ws != null) {
            ws.handleClose();
        }
        wsReqMap.remove(ctx.channel());
    }

    /**
     * Adds an action to be called on HTTP request with
     * {@link ServerHttpExchange}.
     */
    public VibeServerCodec httpAction(Action<ServerHttpExchange> action) {
        httpActions.add(action);
        return this;
    }

    /**
     * Adds an action to be called on WebSocket connection with
     * {@link ServerWebSocket} in open state.
     */
    public VibeServerCodec websocketAction(Action<ServerWebSocket> action) {
        wsActions.add(action);
        return this;
    }

}
@@ -1531,16 +1531,68 @@
 Handler;
+%0Aimport io.netty.handler.timeout.ReadTimeoutHandler;
 %0A%0Aimport
@@ -2222,32 +2222,38 @@
  %7B%0A%0A    private 
+final 
 Actions%3CServerHt
@@ -2313,16 +2313,22 @@
 private 
+final 
 Actions%3C
@@ -2383,32 +2383,38 @@
 ();%0A    private 
+final 
 Map%3CChannel, Net
@@ -2476,32 +2476,38 @@
 ();%0A    private 
+final 
 Map%3CChannel, Net
@@ -2572,16 +2572,22 @@
 private 
+final 
 Map%3CChan
@@ -2979,20 +2979,19 @@
 Headers.
-Valu
+Nam
 es.UPGRA
@@ -3352,16 +3352,136 @@
 quest);%0A
+                // Set timeout to avoid memory leak%0A                ctx.pipeline().addFirst(new ReadTimeoutHandler(5));%0A
         
@@ -3951,32 +3951,32 @@
 tHttpContent) %7B%0A
-
                 
@@ -4003,32 +4003,139 @@
 ctx.channel());%0A
+                    // Cancel timeout%0A                    ctx.pipeline().remove(ReadTimeoutHandler.class);%0A
                 

Set timeout to prevent possible memory leak