/*
 * Copyright @ 2013 CurisTEC, S.A.S. All Rights Reserved.
 */
package net.curisit.securis.utils;

import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.util.zip.GZIPOutputStream;

import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;

/**
 * GZipServletResponseWrapper
 * <p>
 * {@link HttpServletResponseWrapper} that transparently compresses the response
 * body using GZIP. Intended for use in filters/servlets where the caller wants
 * to wrap the original response and write compressed bytes to the client.
 * <p>
 * How it works:
 * <ul>
 *   <li>Overrides {@link #getOutputStream()} and {@link #getWriter()} to route output through a {@link GZIPOutputStream}.</li>
 *   <li>Ensures mutual exclusivity between OutputStream and Writer access per Servlet API requirements.</li>
 *   <li>Ignores {@link #setContentLength(int)} since compressed size differs from uncompressed.</li>
 * </ul>
 *
 * Usage:
 * <pre>
 *   GZipServletResponseWrapper gz = new GZipServletResponseWrapper(resp);
 *   chain.doFilter(request, gz);
 *   gz.close(); // important: finish compression and flush buffers
 * </pre>
 *
 * Thread-safety:
 * <p>
 * Instances are per-request and not shared.
 *
 * Limitations:
 * <p>
 * Caller is responsible for setting "Content-Encoding: gzip" and for avoiding
 * double-compression scenarios.
 *
 * @author JRA
 * Last reviewed by JRA on Oct 6, 2025.
 */
public class GZipServletResponseWrapper extends HttpServletResponseWrapper {

	private GZIPServletOutputStream gzipOutputStream = null;
	private PrintWriter printWriter = null;

	// ---------------------------------------------------------------------
	// Constructors
	// ---------------------------------------------------------------------

	/**
	 * GZipServletResponseWrapper
	 * <p>
	 * Wraps the given response. Actual GZIP streams are lazily created on first write.
	 *
	 * @param response the original {@link HttpServletResponse} to wrap
	 * @throws IOException if the underlying response streams cannot be accessed
	 */
	public GZipServletResponseWrapper(HttpServletResponse response) throws IOException {
		super(response);
	}

	// ---------------------------------------------------------------------
	// Lifecycle
	// ---------------------------------------------------------------------

	/**
	 * close<p>
	 * Closes any open writer or output stream and finalizes the GZIP stream.
	 * Must be called once all response content has been written.
	 *
	 * @throws IOException if closing the underlying streams fails
	 */
	public void close() throws IOException {
		// PrintWriter.close does not throw exceptions. Hence no try-catch block.
		if (this.printWriter != null) {
			this.printWriter.close();
		}
		if (this.gzipOutputStream != null) {
			this.gzipOutputStream.close();
		}
	}

	/**
	 * flushBuffer<p>
	 * Flushes the writer and the GZIP output stream, then delegates to the wrapped response.
	 * If multiple exceptions occur, the first encountered is thrown (typical servlet practice).
	 *
	 * @throws IOException if flushing any of the streams fails
	 */
	@Override
	public void flushBuffer() throws IOException {

		// PrintWriter.flush() does not throw exception
		if (this.printWriter != null) {
			this.printWriter.flush();
		}

		IOException exception1 = null;
		try {
			if (this.gzipOutputStream != null) {
				this.gzipOutputStream.flush();
			}
		} catch (IOException e) {
			exception1 = e;
		}

		IOException exception2 = null;
		try {
			super.flushBuffer();
		} catch (IOException e) {
			exception2 = e;
		}

		if (exception1 != null) throw exception1;
		if (exception2 != null) throw exception2;
	}

	// ---------------------------------------------------------------------
	// Output acquisition
	// ---------------------------------------------------------------------

	/**
	 * getOutputStream<p>
	 * Returns a {@link ServletOutputStream} that writes compressed data.
	 * Mutually exclusive with {@link #getWriter()} as per Servlet API.
	 *
	 * @return compressed {@link ServletOutputStream}
	 * @throws IOException if the underlying output stream cannot be obtained
	 * @throws IllegalStateException if the writer has been already acquired
	 */
	@Override
	public ServletOutputStream getOutputStream() throws IOException {
		if (this.printWriter != null) {
			throw new IllegalStateException("PrintWriter obtained already - cannot get OutputStream");
		}
		if (this.gzipOutputStream == null) {
			this.gzipOutputStream = new GZIPServletOutputStream(getResponse().getOutputStream());
		}
		return this.gzipOutputStream;
	}

	/**
	 * getWriter<p>
	 * Returns a {@link PrintWriter} that writes compressed data (UTF-8 by default, inherited from response).
	 * Mutually exclusive with {@link #getOutputStream()} as per Servlet API.
	 *
	 * @return compressed {@link PrintWriter}
	 * @throws IOException if streams cannot be allocated
	 * @throws IllegalStateException if the output stream has been already acquired
	 */
	@Override
	public PrintWriter getWriter() throws IOException {
		if (this.printWriter == null && this.gzipOutputStream != null) {
			throw new IllegalStateException("OutputStream obtained already - cannot get PrintWriter");
		}
		if (this.printWriter == null) {
			this.gzipOutputStream = new GZIPServletOutputStream(getResponse().getOutputStream());
			this.printWriter = new PrintWriter(new OutputStreamWriter(this.gzipOutputStream, getResponse().getCharacterEncoding()));
		}
		return this.printWriter;
	}

	/**
	 * setContentLength<p>
	 * No-op. The content length of the zipped content is not known a priori and
	 * will not match the uncompressed length; therefore we do not set it here.
	 *
	 * @param len ignored
	 */
	@Override
	public void setContentLength(int len) {
		// ignore, since content length of zipped content does not match content length of unzipped content.
	}

	// ---------------------------------------------------------------------
	// Inner compressed stream
	// ---------------------------------------------------------------------

	/**
	 * GZIPServletOutputStream
	 * <p>
	 * Decorates the original {@link ServletOutputStream} with a {@link GZIPOutputStream}.
	 * Delegates readiness and listener to the underlying (container) stream.
	 * 
	 * @author JRA
     * Last reviewed by JRA on Oct 5, 2025.
	 */
	private static class GZIPServletOutputStream extends ServletOutputStream {
		private final ServletOutputStream servletOutputStream;
		private final GZIPOutputStream gzipStream;

		/**
		 * GZIPServletOutputStream<p>
		 * Creates a new compressed stream wrapper.
		 *
		 * @param servletOutputStream underlying (container-provided) output stream
		 * @throws IOException if the GZIP stream cannot be created
		 */
		public GZIPServletOutputStream(ServletOutputStream servletOutputStream) throws IOException {
			this.servletOutputStream = servletOutputStream;
			this.gzipStream = new GZIPOutputStream(servletOutputStream);
		}

		/** 
		 * isReady<p>
		 * Check if the output stream is ready
		 * {@inheritDoc} 
		 * 
		 * @return isReady
		 */
		@Override
		public boolean isReady() {
			return this.servletOutputStream.isReady();
		}

		/** 
		 * setWriteListener<p>
		 * Set the write listener for the output stream
		 * {@inheritDoc}
		 * 
		 * @param writeListener
		 */
		@Override
		public void setWriteListener(WriteListener writeListener) {
			this.servletOutputStream.setWriteListener(writeListener);
		}

		/** 
		 * write<p>
		 * Write on the gzip stream
		 * {@inheritDoc} 
		 * 
		 * @param b
		 * @throws IOException
		 */
		@Override
		public void write(int b) throws IOException {
			this.gzipStream.write(b);
		}

		/** 
		 * close<p>
		 * Close the gzip stream
		 * {@inheritDoc} 
		 * 
		 * @throws IOException
		 */
		@Override
		public void close() throws IOException {
			this.gzipStream.close();
		}

		/**
		 * flush<p>
		 * Flush the gzip stream
		 * {@inheritDoc} 
		 * 
		 * @throws IOException
		 */
		@Override
		public void flush() throws IOException {
			this.gzipStream.flush();
		}

		/**
		 * finish<p>
		 * Explicitly finishes writing of the GZIP stream, without closing the underlying stream.
		 * Not used by the wrapper but available for completeness.
		 *
		 * @throws IOException if finishing fails
		 */
		@SuppressWarnings("unused")
		public void finish() throws IOException {
			this.gzipStream.finish();
		}
	}
}