package com.example.servlet;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;

import javax.servlet.ServletOutputStream;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;

public class LoggingHttpServletResponse extends HttpServletResponseWrapper {

    final ByteArrayOutputStream baos;

    Stream out;

    PrintWriter writer;

    public LoggingHttpServletResponse(HttpServletResponse response)
            throws IOException {
        super(response);
        baos = new ByteArrayOutputStream();
    }

    public byte[] getBuffer() throws IOException {
        flushBuffer();
        return baos.toByteArray();
    }

    public String getText() throws IOException {
        return new String(getBuffer(), getCharacterEncoding());
    }

    @Override
    public HttpServletResponse getResponse() {
        return (HttpServletResponse) super.getResponse();
    }

    @Override
    public void setResponse(ServletResponse response) {
        throw new UnsupportedOperationException();
    }

    @Override
    public Stream getOutputStream() throws IOException {
        if (out == null) {
            out = new Stream(super.getOutputStream(), baos);
        }
        return out;
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        if (writer == null) {
            writer = new PrintWriter(new OutputStreamWriter(
                    getOutputStream(), getCharacterEncoding()));
        }
        return writer;
    }

    @Override
    public void flushBuffer() throws IOException {
        super.flushBuffer();
        getOutputStream().flush();
    }

    @Override
    public void resetBuffer() {
        super.resetBuffer();
        baos.reset();
    }

    @Override
    public void reset() {
        super.reset();
        baos.reset();
        out = null;
    }

    static class Stream extends ServletOutputStream {

        final ServletOutputStream original;

        final OutputStream internal;

        public Stream(ServletOutputStream original, OutputStream internal) {
            this.original = original;
            this.internal = internal;
        }

        public void write(int b) throws IOException {
            original.write(b);
            internal.write(b);
        }

        @Override
        public void write(byte[] b) throws IOException {
            original.write(b);
            internal.write(b);
        }

        @Override
        public void write(byte[] b, int off, int len) throws IOException {
            original.write(b, off, len);
            internal.write(b, off, len);
        }

        @Override
        public void flush() throws IOException {
            original.flush();
            internal.flush();
        }

        @Override
        public void close() throws IOException {
            original.close();
            internal.close();
        }

    }

}
