package cn.com.poc.common.utils;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;

/**
 * @author alex.yao
 * @date 2025/3/4
 */
public class SSEUtil {

    private final Logger logger = LoggerFactory.getLogger(SSEUtil.class);

    private final ServletOutputStream outputStream;

    public SSEUtil() throws IOException {
        ServletRequestAttributes servletRequestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        HttpServletResponse response = servletRequestAttributes.getResponse();
        if (response == null) {
            throw new IOException("无法获取HttpServletResponse");
        }
        response.setContentType("text/event-stream;charset=UTF-8");
        outputStream = response.getOutputStream();
    }

    public SSEUtil(HttpServletResponse response) throws IOException {
        response.setContentType("text/event-stream;charset=UTF-8");
        outputStream = response.getOutputStream();
    }

    public void send(String data) throws IOException {
        String message = "data: " + data + "\n\n";
        outputStream.write(message.getBytes(StandardCharsets.UTF_8));
        outputStream.flush();
    }

    public boolean complete() {
        try {
            outputStream.close();
        } catch (IOException e) {
            logger.error("无法关闭sse连接:{}", e.getMessage());
            return false;
        }
        return true;
    }

    public boolean completeByError(String errorMsg) {
        try {
            String mess = "data: " + errorMsg + "\n\n";
            outputStream.write(mess.getBytes(StandardCharsets.UTF_8));
            outputStream.flush();
            outputStream.close();
        } catch (IOException e) {
            logger.error("无法关闭sse连接:{}", e.getMessage());
            return false;
        }
        return true;
    }

}
