package com.atlassian.gzipfilter;
   
   import com.atlassian.gzipfilter.util.IOUtils;
   import org.slf4j.Logger;
   import org.slf4j.LoggerFactory;
   
   import java.io.IOException;
   import java.io.OutputStreamWriter;
   import java.io.PrintWriter;
  
  import javax.servlet.ServletOutputStream;
  import javax.servlet.http.HttpServletResponse;
  import javax.servlet.http.HttpServletResponseWrapper;
  
  public class GzipResponseWrapper extends HttpServletResponseWrapper
  {
      private static final Logger log = LoggerFactory.getLogger(GzipResponseWrapper.class);
      protected HttpServletResponse origResponse = null;
      protected ServletOutputStream stream = null;
      protected PrintWriter writer = null;
      private String encoding;
  
      public GzipResponseWrapper(HttpServletResponse response, String encoding)
      {
          super(response);
          this.encoding = encoding;
          origResponse = response;
      }
  
      protected ServletOutputStream createOutputStream() throws IOException
      {
          return new GzipResponseStream(origResponse);
      }
  
      public void finishResponse()
      {
  
          //JRA-29795 If stream was not created - let's create it and close it right after. GzipOutputStream while construction is
          //writing some metadata into stream and since in this place "Content-Encoding:zip" was set and cannot be removed
          //we have to produce valid GZIP stream with appropriate gzip metadata header/footers.
          if(stream == null)
          {
              try
              {
                  stream = createOutputStream();
                  stream.flush();
              }
              catch (IOException e)
              {
                  log.warn("Was unable to create GzipResponseStream. Invalid gzip stream was sent in response body!", e);
              }
          }
          IOUtils.closeQuietly(writer);
          IOUtils.closeQuietly(stream);
      }
  
      public void flushBuffer() throws IOException
      {
          if (stream != null)
              stream.flush();
      }
  
      public ServletOutputStream getOutputStream() throws IOException
      {
          if (writer != null)
              throw new IllegalStateException("getWriter() has already been called!");
  
          if (stream == null)
              stream = createOutputStream();
  
          return stream;
      }
  
      public PrintWriter getWriter() throws IOException
      {
          if (writer != null)
              return writer;
  
          if (stream != null)
              throw new IllegalStateException("getOutputStream() has already been called!");
  
          stream = createOutputStream();
          writer = new PrintWriter(new OutputStreamWriter(stream, encoding));
  
          return writer;
      }
  
      public void setEncoding(String encoding)
      {
          this.encoding = encoding;
      }
  }