package gc.mda.signal_batch.monitoring.filter;

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

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

public class CachedHttpServletResponseWrapper extends HttpServletResponseWrapper {
    
    private final ByteArrayOutputStream cachedContent;
    private final ServletOutputStream outputStream;
    private PrintWriter writer;

    public CachedHttpServletResponseWrapper(HttpServletResponse response) {
        super(response);
        this.cachedContent = new ByteArrayOutputStream();
        this.outputStream = new CachedServletOutputStream(cachedContent);
    }

    @Override
    public ServletOutputStream getOutputStream() {
        return outputStream;
    }

    @Override
    public PrintWriter getWriter() {
        if (writer == null) {
            writer = new PrintWriter(outputStream, true);
        }
        return writer;
    }

    public byte[] getCachedContent() {
        try {
            if (writer != null) {
                writer.flush();
            }
            outputStream.flush();
        } catch (IOException e) {
            // Handle exception
        }
        return cachedContent.toByteArray();
    }

    public void copyContentToResponse() throws IOException {
        byte[] content = getCachedContent();
        if (content.length > 0) {
            getResponse().getOutputStream().write(content);
            getResponse().flushBuffer();
        }
    }

    private static class CachedServletOutputStream extends ServletOutputStream {
        private final ByteArrayOutputStream outputStream;

        public CachedServletOutputStream(ByteArrayOutputStream outputStream) {
            this.outputStream = outputStream;
        }

        @Override
        public void write(int b) {
            outputStream.write(b);
        }

        @Override
        public void write(byte[] b, int off, int len) {
            outputStream.write(b, off, len);
        }

        @Override
        public boolean isReady() {
            return true;
        }

        @Override
        public void setWriteListener(WriteListener writeListener) {
            // Not implemented for this use case
        }
    }
}