package com.huimv.employment.integration.kb;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.util.StringUtils;

import java.nio.charset.StandardCharsets;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 从 {@code /api/console/chat} SSE 字节流中增量提取 assistant 文本，
 * 识别流结束事件，暂存「最后一条终端消息」，并采集 MCP 草稿工具响应证据。
 */
public class ConsoleChatSseAggregator {

    private static final Pattern MCP_DRAFT_HINT = Pattern.compile(
            "MCP__generate_initial_draft|MCP__update_draft",
            Pattern.CASE_INSENSITIVE);

    private static final Pattern DRAFT_ID_PATTERN = Pattern.compile(
            "\"draft_id\"\\s*:\\s*\"?(\\d+)\"?",
            Pattern.CASE_INSENSITIVE);

    private final ObjectMapper objectMapper;
    private final StringBuilder lineBuffer = new StringBuilder();
    private final StringBuilder streamingText = new StringBuilder();
    private final StringBuilder completedReply = new StringBuilder();
    private final StringBuilder mcpDraftEvidence = new StringBuilder();
    private final Set<String> reasoningMessageIds = new HashSet<>();
    private boolean completed;
    /** 暂存的最后一条终端 SSE 行（未转发，待业务层注入卡片字段） */
    private String heldTerminalLine;
    private Long lastMcpDraftId;

    public ConsoleChatSseAggregator(ObjectMapper objectMapper) {
        this.objectMapper = objectMapper;
    }

    public void appendBytes(byte[] buffer, int offset, int length) {
        if (buffer == null || length <= 0) {
            return;
        }
        lineBuffer.append(new String(buffer, offset, length, StandardCharsets.UTF_8));
        drainCompleteLines();
    }

    public void finish() {
        drainCompleteLines();
        if (lineBuffer.length() > 0) {
            consumeSseLine(lineBuffer.toString().trim());
            lineBuffer.setLength(0);
        }
        completed = true;
    }

    public void acceptLine(String line) {
        if (!StringUtils.hasText(line)) {
            return;
        }
        consumeSseLine(line.trim());
    }

    public boolean isCompleted() {
        return completed;
    }

    public String getAssistantText() {
        if (completedReply.length() > 0) {
            return completedReply.toString();
        }
        return streamingText.toString();
    }

    public String getMcpDraftEvidence() {
        return mcpDraftEvidence.toString();
    }

    public Long getLastMcpDraftId() {
        return lastMcpDraftId;
    }

    /**
     * 取出暂存的最后一条终端消息行（取出后清空）。
     */
    public String takeHeldTerminalLine() {
        String line = heldTerminalLine;
        heldTerminalLine = null;
        return line;
    }

    /**
     * 判断 JSON 载荷是否为对话结束/最后一条消息（需 hold 以便注入卡片）。
     */
    public boolean isTerminalPayload(String payload) {
        if (!StringUtils.hasText(payload) || "[DONE]".equals(payload)) {
            return false;
        }
        try {
            JsonNode data = objectMapper.readTree(payload);
            return isTerminalEvent(data);
        } catch (Exception ex) {
            return false;
        }
    }

    /**
     * hold 一条终端行；若已有 hold，返回旧行供立即转发。
     */
    public String holdTerminalLine(String line) {
        String previous = heldTerminalLine;
        heldTerminalLine = line;
        return previous;
    }

    private void drainCompleteLines() {
        int newlineIndex;
        while ((newlineIndex = indexOfLineBreak(lineBuffer)) >= 0) {
            String line = lineBuffer.substring(0, newlineIndex).trim();
            removeThroughLineBreak(lineBuffer, newlineIndex);
            consumeSseLine(line);
        }
    }

    private void consumeSseLine(String line) {
        String payload = normalizeDataPayload(line);
        if (!StringUtils.hasText(payload)) {
            return;
        }
        if ("[DONE]".equals(payload)) {
            completed = true;
            return;
        }
        try {
            JsonNode data = objectMapper.readTree(payload);
            captureMcpDraftEvidence(payload, data);
            consumeJsonEvent(data);
        } catch (Exception ignored) {
            // skip
        }
    }

    /**
     * 采集 MCP 草稿工具调用/响应（通常在 plugin_call / plugin_call_output，不在展示文案中）。
     */
    private void captureMcpDraftEvidence(String payload, JsonNode data) {
        if (!StringUtils.hasText(payload) || !MCP_DRAFT_HINT.matcher(payload).find()) {
            return;
        }
        if (mcpDraftEvidence.length() > 0) {
            mcpDraftEvidence.append('\n');
        }
        mcpDraftEvidence.append(payload);
        Long id = findDraftIdInNode(data);
        if (id == null) {
            id = findDraftIdByRegex(payload);
        }
        if (id != null) {
            lastMcpDraftId = id;
        }
    }

    private Long findDraftIdInNode(JsonNode node) {
        if (node == null || node.isNull()) {
            return null;
        }
        if (node.isObject()) {
            Long direct = asLongId(node.get("draft_id"));
            if (direct == null) {
                direct = asLongId(node.get("draftId"));
            }
            if (direct != null) {
                return direct;
            }
            Long last = null;
            Iterator<Map.Entry<String, JsonNode>> fields = node.fields();
            while (fields.hasNext()) {
                Map.Entry<String, JsonNode> entry = fields.next();
                Long nested = findDraftIdInNode(entry.getValue());
                if (nested != null) {
                    last = nested;
                }
            }
            return last;
        }
        if (node.isArray()) {
            Long last = null;
            for (JsonNode child : node) {
                Long nested = findDraftIdInNode(child);
                if (nested != null) {
                    last = nested;
                }
            }
            return last;
        }
        // plugin 输出有时把 JSON 再包一层字符串
        if (node.isTextual()) {
            String text = node.asText();
            if (!StringUtils.hasText(text) || text.indexOf("draft_id") < 0) {
                return null;
            }
            try {
                return findDraftIdInNode(objectMapper.readTree(text));
            } catch (Exception ex) {
                return findDraftIdByRegex(text);
            }
        }
        return null;
    }

    private static Long findDraftIdByRegex(String text) {
        if (!StringUtils.hasText(text)) {
            return null;
        }
        Matcher matcher = DRAFT_ID_PATTERN.matcher(text);
        Long last = null;
        while (matcher.find()) {
            try {
                long value = Long.parseLong(matcher.group(1));
                if (value > 0) {
                    last = value;
                }
            } catch (NumberFormatException ignored) {
                // skip
            }
        }
        return last;
    }

    private static Long asLongId(JsonNode node) {
        if (node == null || node.isNull()) {
            return null;
        }
        if (node.isNumber()) {
            long value = node.asLong();
            return value > 0 ? value : null;
        }
        if (node.isTextual()) {
            String text = node.asText().trim();
            if (!StringUtils.hasText(text)) {
                return null;
            }
            try {
                long value = Long.parseLong(text);
                return value > 0 ? value : null;
            } catch (NumberFormatException ex) {
                return null;
            }
        }
        return null;
    }

    private void consumeJsonEvent(JsonNode data) {
        if (data == null || data.isNull()) {
            return;
        }
        String objectType = data.path("object").asText("");
        String status = data.path("status").asText("");

        if ("message".equals(objectType) && "reasoning".equals(data.path("type").asText())) {
            String messageId = data.path("id").asText(null);
            if (StringUtils.hasText(messageId)) {
                reasoningMessageIds.add(messageId);
            }
            return;
        }

        if (isTerminalEvent(data)) {
            completed = true;
            if (data.has("output")) {
                String reply = extractAssistantReplyFromOutput(data.get("output"));
                if (StringUtils.hasText(reply)) {
                    completedReply.setLength(0);
                    completedReply.append(reply);
                }
            } else if ("message".equals(objectType) && "assistant".equals(data.path("role").asText())) {
                String reply = extractTextFromContentNodes(data.get("content"));
                if (StringUtils.hasText(reply)) {
                    completedReply.setLength(0);
                    completedReply.append(reply);
                }
            }
            return;
        }

        if ("content".equals(objectType) && "text".equals(data.path("type").asText())) {
            if (data.path("delta").asBoolean(false) && data.has("text")) {
                String msgId = data.path("msg_id").asText(null);
                if (StringUtils.hasText(msgId) && reasoningMessageIds.contains(msgId)) {
                    return;
                }
                streamingText.append(data.get("text").asText(""));
            }
            return;
        }

        JsonNode output = data.get("output");
        if (output != null && output.isArray()) {
            String reply = extractAssistantReplyFromOutput(output);
            if (StringUtils.hasText(reply) && completedReply.length() == 0 && streamingText.length() == 0) {
                streamingText.append(reply);
            }
        }
    }

    private static boolean isTerminalEvent(JsonNode data) {
        if (data == null || !data.isObject()) {
            return false;
        }
        String objectType = data.path("object").asText("");
        String status = data.path("status").asText("");
        String type = data.path("type").asText("");

        if (("response".equals(objectType) || "chat.completion".equals(objectType))
                && ("completed".equals(status) || "complete".equals(status))) {
            return true;
        }
        if ("complete".equals(objectType) || "complete".equals(type) || "completed".equals(type)) {
            return true;
        }
        // 完整 assistant 消息（非 reasoning、非纯 delta）
        if ("message".equals(objectType)
                && "message".equals(data.path("type").asText("message"))
                && "assistant".equals(data.path("role").asText())) {
            return true;
        }
        return false;
    }

    private static String extractAssistantReplyFromOutput(JsonNode output) {
        if (output == null || !output.isArray()) {
            return "";
        }
        StringBuilder reply = new StringBuilder();
        for (JsonNode item : output) {
            if (!"assistant".equals(item.path("role").asText())) {
                continue;
            }
            if (!"message".equals(item.path("type").asText("message"))) {
                continue;
            }
            String text = extractTextFromContentNodes(item.get("content"));
            if (StringUtils.hasText(text)) {
                if (reply.length() > 0) {
                    reply.append('\n');
                }
                reply.append(text);
            }
        }
        return reply.toString();
    }

    private static String extractTextFromContentNodes(JsonNode contents) {
        if (contents == null || !contents.isArray()) {
            return "";
        }
        StringBuilder text = new StringBuilder();
        for (JsonNode content : contents) {
            appendTextPart(text, content);
        }
        return text.toString();
    }

    private static void appendTextPart(StringBuilder target, JsonNode content) {
        if (content == null || content.isNull()) {
            return;
        }
        if (content.has("text") && "text".equals(content.path("type").asText("text"))) {
            target.append(content.get("text").asText(""));
            return;
        }
        if (content.has("text") && content.has("object") && "content".equals(content.path("object").asText())) {
            target.append(content.get("text").asText(""));
        }
    }

    public static String normalizeDataPayload(String line) {
        if (!StringUtils.hasText(line)) {
            return "";
        }
        String payload = line.trim();
        while (payload.startsWith("data:")) {
            payload = payload.substring(5).trim();
        }
        return payload;
    }

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

    private static void removeThroughLineBreak(StringBuilder sb, int newlineIndex) {
        int removeEnd = newlineIndex + 1;
        if (removeEnd < sb.length() && sb.charAt(removeEnd) == '\r') {
            removeEnd++;
        }
        sb.delete(0, removeEnd);
    }
}
