package com.ruoyi.web.kb.support;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import javax.servlet.http.HttpServletResponse;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;

/**
 * Server-Sent Events（SSE）响应头与流式写出工具。
 */
public final class SseStreamSupport
{
    private static final int PIPE_BUFFER_SIZE = 4096;

    private SseStreamSupport()
    {
    }

    /**
     * 设置符合 SSE 规范的响应头（{@code text/event-stream}）。
     */
    public static void prepareSseResponse(HttpServletResponse response)
    {
        response.setStatus(HttpServletResponse.SC_OK);
        response.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE);
        response.setCharacterEncoding(StandardCharsets.UTF_8.name());
        response.setHeader(HttpHeaders.CACHE_CONTROL, "no-cache, no-store, must-revalidate");
        response.setHeader(HttpHeaders.PRAGMA, "no-cache");
        response.setHeader(HttpHeaders.EXPIRES, "0");
        response.setHeader(HttpHeaders.CONNECTION, "keep-alive");
        /** 禁止 Nginx 等反代缓冲，确保客户端实时收到事件 */
        response.setHeader("X-Accel-Buffering", "no");
    }

    /**
     * 将上游 SSE 字节流按块写出并 flush，避免整段缓冲。
     */
    public static void pipeWithFlush(InputStream input, OutputStream output) throws IOException
    {
        if (input == null)
        {
            return;
        }
        byte[] buffer = new byte[PIPE_BUFFER_SIZE];
        int read;
        while ((read = input.read(buffer)) != -1)
        {
            output.write(buffer, 0, read);
            output.flush();
        }
    }

    /**
     * 写入单条 SSE {@code data:} 事件（以空行结束）。
     */
    public static void writeDataEvent(OutputStream output, String dataPayload) throws IOException
    {
        if (dataPayload == null)
        {
            return;
        }
        String line = "data: " + dataPayload + "\n\n";
        output.write(line.getBytes(StandardCharsets.UTF_8));
        output.flush();
    }

    /**
     * 写入 SSE 结束标记（OpenAI 兼容 {@code [DONE]}）。
     */
    public static void writeDoneEvent(OutputStream output) throws IOException
    {
        writeDataEvent(output, "[DONE]");
    }
}
