package com.huimv.employment.integration.kb;

import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;

/**
 * Server-Sent Events（SSE）响应头与流式写出工具。
 */
public final class SseStreamSupport {

    private static final int PIPE_BUFFER_SIZE = 4096;

    private SseStreamSupport() {
    }

    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");
        response.setHeader("X-Accel-Buffering", "no");
    }

    public static void pipeWithFlush(InputStream input, OutputStream output) throws IOException {
        pipeWithFlush(input, output, null);
    }

    public static void pipeWithFlush(InputStream input, OutputStream output, ConsoleChatSseAggregator aggregator)
            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();
            if (aggregator != null) {
                aggregator.appendBytes(buffer, 0, read);
            }
        }
    }

    /**
     * 会话落库：按行转发；暂缓 [DONE] 与「最后一条终端消息」，便于注入卡片字段后再写出。
     */
    public static void pipeCollectingDeferDone(InputStream input,
                                               OutputStream output,
                                               ConsoleChatSseAggregator aggregator) throws IOException {
        if (input == null) {
            return;
        }
        StringBuilder lineBuffer = new StringBuilder();
        byte[] buffer = new byte[PIPE_BUFFER_SIZE];
        int read;
        while ((read = input.read(buffer)) != -1) {
            lineBuffer.append(new String(buffer, 0, read, StandardCharsets.UTF_8));
            drainAndForward(lineBuffer, output, aggregator);
        }
        if (lineBuffer.length() > 0) {
            String leftover = lineBuffer.toString();
            lineBuffer.setLength(0);
            if (leftover.endsWith("\r")) {
                leftover = leftover.substring(0, leftover.length() - 1);
            }
            forwardOrHoldLine(leftover, output, aggregator);
        }
    }

    private static void drainAndForward(StringBuilder lineBuffer,
                                        OutputStream output,
                                        ConsoleChatSseAggregator aggregator) throws IOException {
        int newlineIndex;
        while ((newlineIndex = indexOfLineBreak(lineBuffer)) >= 0) {
            String line = lineBuffer.substring(0, newlineIndex);
            if (line.endsWith("\r")) {
                line = line.substring(0, line.length() - 1);
            }
            lineBuffer.delete(0, newlineIndex + 1);
            forwardOrHoldLine(line, output, aggregator);
        }
    }

    private static void forwardOrHoldLine(String line,
                                          OutputStream output,
                                          ConsoleChatSseAggregator aggregator) throws IOException {
        if (aggregator != null) {
            aggregator.acceptLine(line.trim());
        }
        String payload = ConsoleChatSseAggregator.normalizeDataPayload(line.trim());
        if ("[DONE]".equals(payload)) {
            return;
        }
        // 终端消息：hold，旧 hold 先发出（保持流顺序）
        if (aggregator != null && aggregator.isTerminalPayload(payload)) {
            String previous = aggregator.holdTerminalLine(line);
            if (previous != null) {
                writeLine(output, previous);
            }
            return;
        }
        writeLine(output, line);
    }

    private static void writeLine(OutputStream output, String line) throws IOException {
        output.write((line + "\n").getBytes(StandardCharsets.UTF_8));
        output.flush();
    }

    private static int indexOfLineBreak(StringBuilder sb) {
        for (int i = 0; i < sb.length(); i++) {
            if (sb.charAt(i) == '\n') {
                return i;
            }
        }
        return -1;
    }

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

    public static void writeDoneEvent(OutputStream output) throws IOException {
        writeDataEvent(output, "[DONE]");
    }

    /** 原样写出 hold 的 SSE 行（补齐事件分隔）。 */
    public static void writeRawSseLine(OutputStream output, String line) throws IOException {
        if (output == null || line == null) {
            return;
        }
        String trimmed = line.trim();
        if (trimmed.isEmpty()) {
            return;
        }
        if (!trimmed.startsWith("data:")) {
            trimmed = "data: " + trimmed;
        }
        output.write((trimmed + "\n\n").getBytes(StandardCharsets.UTF_8));
        output.flush();
    }
}
