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();
 
            // Reuse the original handshake request but replace the body.
            final Request request = parseRequest(ctx, nettyRequest);
            ChannelBuffer b = nettyWebSocketFrame.getBinaryData();
            // Use params.body in Play-Controller to retrieve the binary content.
            request.body = new ChannelBufferInputStream(b);
            // Mark this request as using Websocket method
            request.method = WEBSOCKET_METHOD_NAME;
            // Also provide channel reference in request object for use in PlayController
            request.args.put("CHANNEL", ctx.getChannel());
 
            final Response response = new Response();
 
            try {
                Http.Response.current.set(response);
                response.out = new ByteArrayOutputStream();
 
                Invoker.invoke(new NettyInvocation(request, response, ctx, nettyRequest, e));
 
            } catch (Exception ex) {
                serve500(ex, ctx, nettyRequest);
            }
 
        } else {
            // Not a websocket frame, pass to PlayHandler.
            // As there is no way to override an inner class in java, a simple workaround is used here.
            // super.messageReceived(ctx, e);
            superClassMessageReceived(ctx, e);
        }
 
        Logger.trace("messageReceived: end");
    }
 
    public class NettyInvocation extends PlayHandler.NettyInvocation {
 
        private final ChannelHandlerContext ctx;
        private final Request request;
        private final Response response;
        private final HttpRequest nettyRequest;
        private final MessageEvent event;
        private final boolean isWebsocketUpgradeRequest;
 
        public NettyInvocation(Request request, Response response, ChannelHandlerContext ctx, HttpRequest nettyRequest, MessageEvent e) {
            super(request, response, ctx, nettyRequest, e);
            this.ctx = ctx;
            this.request = request;
            this.response = response;
            this.nettyRequest = nettyRequest;
            this.event = e;
 
            this.isWebsocketUpgradeRequest = HttpHeaders.Values.UPGRADE.equalsIgnoreCase(nettyRequest.getHeader(HttpHeaders.Names.CONNECTION)) &&
                HttpHeaders.Values.WEBSOCKET.equalsIgnoreCase(nettyRequest.getHeader(HttpHeaders.Names.UPGRADE)) &&
                !WEBSOCKET_METHOD_NAME.equalsIgnoreCase(request.method);
        }
 
        // !!!!!! We want this to be invoked by the original PlayHandler Class
        @Override
        public boolean init() {
 
            if (isWebsocketUpgradeRequest) {
                Logger.trace("init: begin");
                Logger.trace("init: WebsocketUpgrade deferred to execute()");
                Logger.trace("init: end");
                return true;  // Do not prevent execute()
 
            } else {
                // Not a websocket upgrade request, pass to PlayHandler.
                return super.init();
            }
        }
 
        @Override
        public void execute() throws Exception {
            Logger.trace("execute: begin");
 
            if (isWebsocketUpgradeRequest) {
                // This block is to do deferred handshake.
                handleWebSocketHandshakeRequest(ctx, nettyRequest);
                // Keep nettyRequest reference in Map.
                activeWebsocketMap.put(ctx, nettyRequest);
                // Faster version of the above
                //ctx.setAttachment(nettyRequest);
 
            } else if (WEBSOCKET_METHOD_NAME.equalsIgnoreCase(request.method)) {
                // Only do this block for replies to upgraded websocket streams.
                ActionInvoker.invoke(request, response);
                ctx.getChannel().write(new DefaultWebSocketFrame(response.out.toString()));
 
            } else {
                super.execute();
            }
 
            Logger.trace("execute: 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());
    }
 
    /*
     * This method is an exact copy of PlayHandler.messageReceived()
     * This is a simple workaround. You do not need to look into it.
     */
    public void superClassMessageReceived(ChannelHandlerContext ctx, MessageEvent e) throws Exception {
        Logger.trace("messageReceived: begin");
 
        final Object msg = e.getMessage();
        if (msg instanceof HttpRequest) {
            final HttpRequest nettyRequest = (HttpRequest) msg;
            try {
                Request request = parseRequest(ctx, nettyRequest);
                request = processRequest(request);
 
                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);
            }
        }
        Logger.trace("messageReceived: end");
    }
 
}