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;

public class GZipServletResponseWrapper extends HttpServletResponseWrapper {

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

	public GZipServletResponseWrapper(HttpServletResponse response) throws IOException {
		super(response);
	}

	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();
		}
	}

	/**
	 * Flush OutputStream or PrintWriter
	 *
	 * @throws IOException
	 */

	@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;
	}

	@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;
	}

	@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;
	}

	@Override
	public void setContentLength(int len) {
		//ignore, since content length of zipped content
		//does not match content length of unzipped content.
	}

	private static class GZIPServletOutputStream extends ServletOutputStream {
		private final ServletOutputStream servletOutputStream;
		private final GZIPOutputStream gzipStream;

		public GZIPServletOutputStream(ServletOutputStream servletOutputStream) throws IOException {
			this.servletOutputStream = servletOutputStream;
			this.gzipStream = new GZIPOutputStream(servletOutputStream);
		}

		@Override
		public boolean isReady() {
			return this.servletOutputStream.isReady();
		}

		@Override
		public void setWriteListener(WriteListener writeListener) {
			this.servletOutputStream.setWriteListener(writeListener);
		}

		@Override
		public void write(int b) throws IOException {
			this.gzipStream.write(b);
		}

		@Override
		public void close() throws IOException {
			this.gzipStream.close();
		}

		@Override
		public void flush() throws IOException {
			this.gzipStream.flush();
		}

		@SuppressWarnings("unused")
		public void finish() throws IOException {
			this.gzipStream.finish();
		}
	}

}