package play.server;
 
import org.jboss.netty.buffer.ChannelBuffer;
import org.jboss.netty.buffer.ChannelBufferInputStream;
import org.jboss.netty.buffer.ChannelBuffers;
import org.jboss.netty.channel.*;
import org.jboss.netty.handler.codec.http.*;
import org.jboss.netty.handler.codec.http.websocket.*;
 
import static org.jboss.netty.handler.codec.http.HttpHeaders.Names.*;
 
import play.Invoker;
import play.Logger;
import play.Play;
import play.PlayPlugin;
import play.mvc.ActionInvoker;
import play.mvc.Http;
import play.mvc.Http.Request;
import play.mvc.Http.Response;
 
import java.io.*;
import java.util.logging.Level;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.security.MessageDigest;
 
public class PlayWebsocketHandler extends PlayHandler {
 
    // Specify ws:// or wss:// here
    private static final String WEBSOCKET_SCHEME = "ws://";
 
    // Specify PlayRouter method prefix here
    private static final String WEBSOCKET_METHOD_NAME = "WEBSOCKET";
 
    /* A map that stores netty-requests that has been upgraded (via handshake) to websocket stream.
     * ConcurrentMap provides good lookup performance.
     * ChannelHandlerContext can be used as reference, as it does not change on every messageReceived.
     * Refer to http://docs.jboss.org/netty/3.1/api/org/jboss/netty/channel/ChannelHandlerContext.html#setAttachment
     * Alternatively, client ip address may be used as key.
     */
    private final ConcurrentMap<ChannelHandlerContext, HttpRequest> activeWebsocketMap = new ConcurrentSkipListMap<ChannelHandlerContext, HttpRequest>();
 
    @Override
    public void messageReceived(ChannelHandlerContext ctx, MessageEvent e) throws Exception {
        Logger.trace("messageReceived: begin");
 
        final Object msg = e.getMessage();
        if (msg instanceof WebSocketFrame) {
            final WebSocketFrame nettyWebSocketFrame = (WebSocketFrame) msg;
 
            // This is all that is needed to send arbitrary content over the wire.
            // ctx.getChannel().write(new DefaultWebSocketFrame(nettyWebSocketFrame.getTextData().toUpperCase()));
 
            // Websocket frames do not give netty requests, and thus need to be retrieved from the handshake.
            HttpRequest nettyRequest = activeWebsocketMap.get(ctx);
            // Faster version of the above
            //HttpRequest nettyRequest = (HttpRequest) ctx.getAttachment();
 
            try {
                // Needed so parseRequest() will work properly
                nettyRequest.setHeader("X-HTTP-Method-Override", WEBSOCKET_METHOD_NAME);
                ChannelBuffer b = nettyWebSocketFrame.getBinaryData();
                nettyRequest.setContent(b);
 
                // Should optimize to avoid doing parsing on every frame
                Request request = parseRequest(ctx, nettyRequest);
                request = processRequest(request);
 
                // More optimized version than doing setContent() for parseRequest()
                request.body = new ChannelBufferInputStream(b);
                //
 
                final Response response = new Response();
 
                Http.Response.current.set(response);
                response.out = new ByteArrayOutputStream();
                boolean raw = false;
                for (PlayPlugin plugin : Play.plugins) {
                    if (plugin.rawInvocation(request, response)) {
                        raw = true;
                        break;
                    }
                }
                if (raw) {
                    copyResponse(ctx, request, response, nettyRequest);
                } else {
                    Invoker.invoke(new NettyInvocation(request, response, ctx, nettyRequest, e));
                }
 
            } catch (Exception ex) {
                serve500(ex, ctx, nettyRequest);
            }
 
        } else if (msg instanceof HttpRequest) {
            final HttpRequest nettyRequest = (HttpRequest) msg;
 
            if (HttpHeaders.Values.UPGRADE.equalsIgnoreCase(nettyRequest.getHeader(HttpHeaders.Names.CONNECTION)) &&
                HttpHeaders.Values.WEBSOCKET.equalsIgnoreCase(nettyRequest.getHeader(HttpHeaders.Names.UPGRADE))) {
 
                // Do upgrade handshake.
                handleWebSocketHandshakeRequest(ctx, nettyRequest);
                // Keep nettyRequest reference in Map.
                activeWebsocketMap.put(ctx, nettyRequest);
                // Faster version of the above
                //ctx.setAttachment(nettyRequest);
 
            }
 
        } else {
            // Not a websocket frame nor upgrade request, pass to PlayHandler.
            super.messageReceived(ctx, e);
        }
 
        Logger.trace("messageReceived: end");
    }
 
    /*
     * Returns the websocket endpoint's Uniform Resource Identifier (URI)
     */
    private static String getWebSocketURI(HttpRequest nettyRequest) {
        return WEBSOCKET_SCHEME + nettyRequest.getHeader(HttpHeaders.Names.HOST) + nettyRequest.getUri();
    }
 
    /*
     * This method is an almost exact copy of Netty's sample-code for Websocket.
     * This is a proven implementation. You do not need to look into it.
     */
    private static void handleWebSocketHandshakeRequest(ChannelHandlerContext ctx, HttpRequest nettyRequest) {
        // Create the WebSocket handshake response.
        HttpResponse res = new DefaultHttpResponse(
                HttpVersion.HTTP_1_1,
                new HttpResponseStatus(101, "Web Socket Protocol Handshake"));
        res.addHeader(HttpHeaders.Names.UPGRADE, HttpHeaders.Values.WEBSOCKET);
        res.addHeader(HttpHeaders.Names.CONNECTION, HttpHeaders.Values.UPGRADE);
 
        // Fill in the headers and contents depending on handshake method.
        if (nettyRequest.containsHeader(SEC_WEBSOCKET_KEY1) &&
            nettyRequest.containsHeader(SEC_WEBSOCKET_KEY2)) {
            // New handshake method with a challenge:
            res.addHeader(SEC_WEBSOCKET_ORIGIN, nettyRequest.getHeader(ORIGIN));
            res.addHeader(SEC_WEBSOCKET_LOCATION, getWebSocketURI(nettyRequest));
            String protocol = nettyRequest.getHeader(SEC_WEBSOCKET_PROTOCOL);
            if (protocol != null) {
                res.addHeader(SEC_WEBSOCKET_PROTOCOL, protocol);
            }
 
            // Calculate the answer of the challenge.
            String key1 = nettyRequest.getHeader(SEC_WEBSOCKET_KEY1);
            String key2 = nettyRequest.getHeader(SEC_WEBSOCKET_KEY2);
            int a = (int) (Long.parseLong(key1.replaceAll("[^0-9]", "")) / key1.replaceAll("[^ ]", "").length());
            int b = (int) (Long.parseLong(key2.replaceAll("[^0-9]", "")) / key2.replaceAll("[^ ]", "").length());
            long c = nettyRequest.getContent().readLong();
            ChannelBuffer input = ChannelBuffers.buffer(16);
            input.writeInt(a);
            input.writeInt(b);
            input.writeLong(c);
            byte[] digest = null;
            try {
                digest = MessageDigest.getInstance("MD5").digest(input.array());
            } catch (Exception ex) {
                java.util.logging.Logger.getLogger(PlayHandler.class.getName()).log(Level.SEVERE, null, ex);
            }
            ChannelBuffer output = ChannelBuffers.wrappedBuffer(digest);
            res.setContent(output);
        } else {
            // Old handshake method with no challenge:
            res.addHeader(WEBSOCKET_ORIGIN, nettyRequest.getHeader(ORIGIN));
            res.addHeader(WEBSOCKET_LOCATION, getWebSocketURI(nettyRequest));
            String protocol = nettyRequest.getHeader(WEBSOCKET_PROTOCOL);
            if (protocol != null) {
                res.addHeader(WEBSOCKET_PROTOCOL, protocol);
            }
        }
 
        // Upgrade the connection and send the handshake response.
        ChannelPipeline p = ctx.getChannel().getPipeline();
        p.remove("aggregator");
        p.replace("decoder", "wsdecoder", new WebSocketFrameDecoder());
 
        ctx.getChannel().write(res);
 
        p.replace("encoder", "wsencoder", new WebSocketFrameEncoder());
    }
 
}