package com.ruoyi.common.utils;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;

/**
 * 响应包装器 - 用于修改响应内容
 */
public class EncryptHttpServletResponseWrapper extends HttpServletResponseWrapper {
    private ByteArrayOutputStream buffer;
    private ServletOutputStream out;
    private PrintWriter writer;
    
    public EncryptHttpServletResponseWrapper(HttpServletResponse response) {
        super(response);
        buffer = new ByteArrayOutputStream();
        out = new WrappedOutputStream(buffer);
        writer = new PrintWriter(new OutputStreamWriter(buffer, StandardCharsets.UTF_8));
    }
    
    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return out;
    }
    
    @Override
    public PrintWriter getWriter() throws IOException {
        return writer;
    }
    
    @Override
    public void flushBuffer() throws IOException {
        if (out != null) {
            out.flush();
        }
        if (writer != null) {
            writer.flush();
        }
    }
    
    public byte[] getContent() throws IOException {
        flushBuffer();
        return buffer.toByteArray();
    }
    
    private static class WrappedOutputStream extends ServletOutputStream {
        private ByteArrayOutputStream buffer;
        
        public WrappedOutputStream(ByteArrayOutputStream buffer) {
            this.buffer = buffer;
        }
        
        @Override
        public void write(int b) throws IOException {
            buffer.write(b);
        }
        
        @Override
        public boolean isReady() {
            return false;
        }
        
        @Override
        public void setWriteListener(WriteListener writeListener) {
        }
    }
}