package org.webbitserver.netty;

import org.jboss.netty.buffer.ChannelBuffer;
import org.jboss.netty.buffer.ChannelBuffers;
import org.jboss.netty.channel.ChannelFuture;
import org.jboss.netty.channel.ChannelFutureListener;
import org.jboss.netty.channel.ChannelHandlerContext;
import org.jboss.netty.handler.codec.frame.TooLongFrameException;
import org.jboss.netty.handler.codec.http.HttpHeaders;
import org.jboss.netty.handler.codec.http.HttpResponse;
import org.jboss.netty.handler.codec.http.HttpResponseStatus;
import org.jboss.netty.handler.codec.http.HttpHeaders.Names;
import org.jboss.netty.handler.codec.http.HttpHeaders.Values;
import org.jboss.netty.util.CharsetUtil;
import org.webbitserver.WebbitException;
import org.webbitserver.helpers.DateHelper;
import org.jboss.netty.handler.codec.http.CookieEncoder;
import org.jboss.netty.handler.codec.http.DefaultCookie;
import org.jboss.netty.handler.codec.http.Cookie;
import org.jboss.netty.handler.codec.http.DefaultHttpChunk;

import java.io.PrintWriter;
import java.io.StringWriter;
import java.net.HttpCookie;
import java.nio.ByteBuffer;
import java.nio.charset.Charset;
import java.util.Date;

import static org.jboss.netty.buffer.ChannelBuffers.copiedBuffer;
import static org.jboss.netty.buffer.ChannelBuffers.wrappedBuffer;

public class NettyHttpResponse implements org.webbitserver.HttpResponse {

    private static final Charset DEFAULT_CHARSET = Charset.forName("UTF-8");
    private final ChannelHandlerContext ctx;
    private final HttpResponse response;
    private final boolean isKeepAlive;
    private final Thread.UncaughtExceptionHandler exceptionHandler;
    private final ChannelBuffer responseBuffer;
    private Charset charset;

    public NettyHttpResponse(ChannelHandlerContext ctx,
                             HttpResponse response,
                             boolean isKeepAlive,
                             Thread.UncaughtExceptionHandler exceptionHandler) {
        this.ctx = ctx;
        this.response = response;
        this.isKeepAlive = isKeepAlive;
        this.exceptionHandler = exceptionHandler;
        this.charset = DEFAULT_CHARSET;
        responseBuffer = ChannelBuffers.dynamicBuffer();
    }

    @Override
    public NettyHttpResponse charset(Charset charset) {
        this.charset = charset;
        return this;
    }

    @Override
    public Charset charset() {
        return charset;
    }

    @Override
    public NettyHttpResponse status(int status) {
        response.setStatus(HttpResponseStatus.valueOf(status));
        return this;
    }

    @Override
    public NettyHttpResponse chunked() {
        response.setHeader(Names.TRANSFER_ENCODING, Values.CHUNKED);
        response.setChunked(true);
        ctx.getChannel().write(response);
        return this;
    }

    @Override
    public int status() {
        return response.getStatus().getCode();
    }

    @Override
    public NettyHttpResponse header(String name, String value) {
        if (value == null) {
            response.removeHeader(name);
        } else {
            response.addHeader(name, value);
        }
        return this;
    }

    @Override
    public NettyHttpResponse header(String name, long value) {
        response.addHeader(name, value);
        return this;
    }

    @Override
    public NettyHttpResponse header(String name, Date value) {
        response.addHeader(name, DateHelper.rfc1123Format(value));
        return this;
    }

    @Override
    public boolean containsHeader(String name) {
        return response.containsHeader(name);
    }

    @Override
    public NettyHttpResponse cookie(HttpCookie httpCookie) {
        Cookie nettyCookie = new DefaultCookie(httpCookie.getName(),httpCookie.getValue());
        nettyCookie.setDomain(httpCookie.getDomain());
        nettyCookie.setPath(httpCookie.getPath());
        nettyCookie.setSecure(httpCookie.getSecure());
        nettyCookie.setMaxAge((int)httpCookie.getMaxAge());
        nettyCookie.setVersion(httpCookie.getVersion());
        nettyCookie.setDiscard(httpCookie.getDiscard());
        nettyCookie.setHttpOnly(true);
        CookieEncoder encoder = new CookieEncoder(true);
        encoder.addCookie(nettyCookie);
        return header(HttpHeaders.Names.SET_COOKIE, encoder.encode());
    }

    @Override
    public NettyHttpResponse content(String content) {
        return content(copiedBuffer(content, charset()));
    }

    @Override
    public NettyHttpResponse content(byte[] content) {
        return content(copiedBuffer(content));
    }

    @Override
    public NettyHttpResponse content(ByteBuffer buffer) {
        return content(wrappedBuffer(buffer));
    }

    private NettyHttpResponse content(ChannelBuffer content) {
        if (response.isChunked()) {
            throw new UnsupportedOperationException();
        }
        responseBuffer.writeBytes(content);
        return this;
    }

    @Override
    public NettyHttpResponse write(String content) {
        if (response.isChunked()) {
            ctx.getChannel().write(new DefaultHttpChunk(wrappedBuffer(content.getBytes(CharsetUtil.UTF_8))));  
        } else {
            write(copiedBuffer(content, CharsetUtil.UTF_8));
        }    
        return this;
    }

    @Override
    public NettyHttpResponse error(Throwable error) {
        if (error instanceof TooLongFrameException) {
            response.setStatus(HttpResponseStatus.REQUEST_ENTITY_TOO_LARGE);
        } else {
            response.setStatus(HttpResponseStatus.INTERNAL_SERVER_ERROR);
        }
        String message = getStackTrace(error);
        header("Content-Type", "text/plain");
        content(message);
        flushResponse();

        exceptionHandler.uncaughtException(Thread.currentThread(),
                WebbitException.fromException(error, ctx.getChannel()));

        return this;
    }

    private String getStackTrace(Throwable error) {
        StringWriter buffer = new StringWriter();
        PrintWriter writer = new PrintWriter(buffer);
        error.printStackTrace(writer);
        writer.flush();
        return buffer.toString();
    }

    @Override
    public NettyHttpResponse end() {
        flushResponse();
        return this;
    }

    private void flushResponse() {
        try {
            // TODO: Shouldn't have to do this, but without it we sometimes seem to get two Content-Length headers in the response.
            header("Content-Length", (String) null);
            header("Content-Length", responseBuffer.readableBytes());
            ChannelFuture  future = response.isChunked() ? ctx.getChannel().write(new DefaultHttpChunk(ChannelBuffers.EMPTY_BUFFER)) : write(responseBuffer);
            if (!isKeepAlive) {
                future.addListener(ChannelFutureListener.CLOSE);
            }
        } catch (Exception e) {
            exceptionHandler.uncaughtException(Thread.currentThread(),
                    WebbitException.fromException(e, ctx.getChannel()));
        }
    }

    private ChannelFuture write(ChannelBuffer responseBuffer) {
        response.setContent(responseBuffer);
        return ctx.getChannel().write(response);
    }

}

@@ -5667,33 +5667,187 @@
         
-flushResponse();%0A
+try%7B%0A            flushResponse();%0A        %7Dcatch (IllegalStateException e)%7B%0A            return null;%0A        %7Dcatch (WebbitException e)%7B%0A            return null;%0A        %7D
 %0A       

Update NettyHttpResponse.java