RAG/src/main/java/com/wok/supportbot/app/SuggestionResponseParser.java

package com.wok.supportbot.app;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.wok.supportbot.cache.SuggestionCache;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 建议问题解析器：从 LLM 原始输出中分离 answer 和 suggestions。
 * <p>
 * 哨兵标记为 <code>___SUGGESTIONS___</code>，之后为 JSON 字符串数组。
 * 支持同步路径（直接分割完整文本）和流式路径（滑动窗口检测哨兵）。
 */
@Slf4j
public final class SuggestionResponseParser {

    /** 主哨兵：严格匹配提示词要求的下划线分隔标记 */
    private static final String SENTINEL = "___SUGGESTIONS___";

    /** 兼容哨兵：模型偶尔不按指令输出时的大写变形 */
    private static final String LOOSE_SENTINEL = "SUGGESTIONS";

    private static final ObjectMapper objectMapper = new ObjectMapper();

    private SuggestionResponseParser() {
    }

    /**
     * 查找最佳哨兵位置：优先严格哨兵，其次兼容哨兵。
     * <p>
     * 兼容哨兵必须后面紧跟 JSON 数组特征（'[' 或空白 + '['），
     * 避免正文中出现 "SUGGESTIONS" 普通单词时误触发。
     *
     * @param text 要搜索的文本
     * @return 哨兵起始位置，未找到返回 -1
     */
    private static int findSentinelIndex(String text) {
        int strictIdx = text.lastIndexOf(SENTINEL);
        if (strictIdx >= 0) {
            return strictIdx;
        }

        // 兼容模式：从后往前找大写 SUGGESTIONS，且后续需连接 JSON 数组
        int idx = text.lastIndexOf(LOOSE_SENTINEL);
        while (idx >= 0) {
            int after = idx + LOOSE_SENTINEL.length();
            if (after < text.length()) {
                char c = text.charAt(after);
                // 允许 SUGGESTIONS[...]、SUGGESTIONS [...]、SUGGESTIONS:\n[...]
                if (c == '[' || Character.isWhitespace(c) || c == ':' || c == '-') {
                    return idx;
                }
            }
            // 继续向前查找更早的兼容哨兵
            idx = text.lastIndexOf(LOOSE_SENTINEL, idx - 1);
        }
        return -1;
    }

    /**
     * 同步路径：按最后一次出现哨兵的位置分割完整文本。
     *
     * @param rawText LLM 原始输出
     * @return 解析结果（answer + suggestions）
     */
    public static ParsedResponse parse(String rawText) {
        if (rawText == null || rawText.isEmpty()) {
            return new ParsedResponse(rawText != null ? rawText : "", Collections.emptyList());
        }

        int lastIdx = findSentinelIndex(rawText);
        if (lastIdx < 0) {
            // 无哨兵标记，整段作为 answer
            return new ParsedResponse(rawText, Collections.emptyList());
        }

        String answer = rawText.substring(0, lastIdx).trim();
        // 按实际匹配的哨兵长度截取后续内容（严格哨兵或兼容哨兵长度不同）
        int matchedSentinelLength = rawText.startsWith(SENTINEL, lastIdx) ? SENTINEL.length() : LOOSE_SENTINEL.length();
        String suggestionsPart = rawText.substring(lastIdx + matchedSentinelLength).trim();

        List<String> suggestions = parseSuggestionsJson(suggestionsPart);
        return new ParsedResponse(answer, suggestions);
    }

    /**
     * 流式路径：从原始 Flux 中分离 answer 和 suggestions。
     * <p>
     * 采用「安全缓冲区」策略防止哨兵泄漏到 UI：
     * <ol>
     *   <li>每次 chunk 追加到累积缓冲区</li>
     *   <li>安全区域 = 缓冲区去掉末尾 SENTINEL.length() 字符（预留哨兵跨 chunk 截断空间）</li>
     *   <li>在安全区域内检查哨兵：找到则发出哨兵前的剩余内容，后续全进 suggestionsBuffer</li>
     *   <li>未找到哨兵则发出安全区域内容</li>
     *   <li>doOnComplete 时：哨兵已发现则解析 suggestions；未发现则发出安全缓冲区内剩余内容</li>
     * </ol>
     *
     * @param rawStream LLM 原始输出流
     * @param cache     建议缓存
     * @param chatId    会话 ID
     * @return 纯 answer 的 Flux（不含哨兵及之后的 suggestions）
     */
    /**
     * 可变状态容器（供 lambda 内部修改）。
     */
    private static class StreamState {
        final StringBuilder buf = new StringBuilder();
        int emitted = 0;
        boolean sentinelFound = false;
    }

    public static Flux<String> parseFromStream(Flux<String> rawStream, SuggestionCache cache, String chatId) {
        StreamState state = new StreamState();

        return rawStream
                .concatMap(chunk -> Flux.<String>create(sink -> {
                    if (state.sentinelFound) {
                        sink.complete();
                        return;
                    }
                    state.buf.append(chunk);

                    int maxSentinelLen = Math.max(SENTINEL.length(), LOOSE_SENTINEL.length());
                    int safeEnd = Math.max(0, state.buf.length() - maxSentinelLen);
                    if (safeEnd <= state.emitted) {
                        sink.complete();
                        return;
                    }

                    String safeRegion = state.buf.substring(state.emitted, safeEnd);
                    int sentinelIdx = findSentinelIndex(safeRegion);

                    if (sentinelIdx >= 0) {
                        state.sentinelFound = true;
                        String beforeSentinel = safeRegion.substring(0, sentinelIdx);
                        if (!beforeSentinel.isEmpty()) {
                            sink.next(beforeSentinel);
                        }
                        // 确定实际匹配到的哨兵长度
                        int matchedLen = safeRegion.startsWith(SENTINEL, sentinelIdx)
                                ? SENTINEL.length()
                                : LOOSE_SENTINEL.length();
                        // 哨兵在安全区域内的位置 sentinelIdx，相对全缓冲区即 emitted + sentinelIdx
                        // emitted 推进到哨兵结束后，后续内容进 suggestions 解析
                        state.emitted = state.emitted + sentinelIdx + matchedLen;
                        log.debug("流式路径检测到哨兵: chatId={}, pos={}, matchedLen={}", chatId, state.emitted, matchedLen);
                    } else {
                        if (!safeRegion.isEmpty()) {
                            sink.next(safeRegion);
                        }
                        state.emitted = safeEnd;
                    }
                    sink.complete();
                }))
                // 流结束后发出安全缓冲区内未发出的残留内容（无哨兵场景）
                .concatWith(Flux.defer(() -> {
                    if (!state.sentinelFound && state.emitted < state.buf.length()) {
                        String residual = state.buf.substring(state.emitted);
                        if (!residual.isEmpty()) {
                            return Flux.just(residual);
                        }
                    }
                    return Flux.empty();
                }))
                .doOnComplete(() -> {
                    if (state.sentinelFound) {
                        String suggestionsPart = state.buf.length() > state.emitted
                                ? state.buf.substring(state.emitted) : "";
                        List<String> suggestions = parseSuggestionsJson(suggestionsPart.strip());
                        if (!suggestions.isEmpty()) {
                            cache.put(chatId, suggestions);
                            log.info("流式 suggestions 解析成功: chatId={}, count={}", chatId, suggestions.size());
                        } else {
                            log.debug("流式 suggestions 解析为空: chatId={}", chatId);
                        }
                    }
                });
    }

    /**
     * 容错解析 suggestions JSON 数组。
     * 先尝试标准 JSON 解析，失败则按行切分取前 3 条非空行。
     */
    static List<String> parseSuggestionsJson(String jsonPart) {
        if (jsonPart == null || jsonPart.isBlank()) {
            return Collections.emptyList();
        }

        // 1. 尝试标准 JSON 解析
        String trimmed = jsonPart.strip();
        // 剥离旧版哨兵标记（___SUGGESTIONS___ 及兼容大写变形），取哨兵后的 JSON 部分。
        // 旧版 suggestion_prompt 会要求模型输出哨兵，若未剥离会被按行降级解析成垃圾条目。
        int sentinelIdx = findSentinelIndex(trimmed);
        if (sentinelIdx >= 0) {
            int matchedLen = trimmed.startsWith(SENTINEL, sentinelIdx)
                    ? SENTINEL.length()
                    : LOOSE_SENTINEL.length();
            trimmed = trimmed.substring(sentinelIdx + matchedLen).strip();
        }
        // 去掉可能的 markdown 代码块包裹
        trimmed = trimCodeBlock(trimmed);

        try {
            List<String> parsed = objectMapper.readValue(trimmed, new TypeReference<List<String>>() {});
            List<String> result = new ArrayList<>();
            for (String s : parsed) {
                if (s != null && !s.isBlank()) {
                    result.add(s.strip());
                }
            }
            if (!result.isEmpty()) {
                return result.size() <= 3 ? result : result.subList(0, 3);
            }
        } catch (Exception e) {
            log.debug("标准 JSON 解析 suggestions 失败，尝试按行切分: {}", e.getMessage());
        }

        // 2. 降级：按行切分，取前 3 条
        return parseByLines(trimmed);
    }

    /**
     * 降级解析：按行切分，去掉序号前缀后取前 3 条非空行。
     */
    private static List<String> parseByLines(String text) {
        List<String> lines = text.lines()
                .map(String::strip)
                .map(SuggestionResponseParser::stripNumberPrefix)
                .map(SuggestionResponseParser::stripQuotes)
                .filter(s -> !s.isBlank())
                .collect(Collectors.toList());

        if (lines.isEmpty()) {
            return Collections.emptyList();
        }
        return lines.size() <= 3 ? lines : lines.subList(0, 3);
    }

    /** 去掉可能的 markdown 代码块包裹（支持 3 个及以上反引号） */
    private static String trimCodeBlock(String s) {
        if (s.startsWith("```") && s.endsWith("```")) {
            // 计算开头的反引号数量
            int openCount = 0;
            while (openCount < s.length() && s.charAt(openCount) == '`') openCount++;
            if (openCount >= 3 && s.endsWith("`".repeat(openCount))) {
                String inner = s.substring(openCount, s.length() - openCount).strip();
                if (inner.startsWith("json")) {
                    inner = inner.substring(4).strip();
                } else if (inner.startsWith("JSON")) {
                    inner = inner.substring(4).strip();
                }
                return inner;
            }
        }
        return s;
    }

    /** 去掉行首序号，如 "1." "2." "3." "1、" "2、" "3、" */
    private static String stripNumberPrefix(String s) {
        if (s.length() >= 2 && Character.isDigit(s.charAt(0)) &&
            (s.charAt(1) == '.' || s.charAt(1) == '、' || s.charAt(1) == ')')) {
            return s.substring(2).strip();
        }
        return s;
    }

    /** 去掉首尾引号 */
    private static String stripQuotes(String s) {
        if (s.length() >= 2) {
            char first = s.charAt(0);
            char last = s.charAt(s.length() - 1);
            if ((first == '"' && last == '"') || (first == '\'' && last == '\'') ||
                (first == '“' && last == '”')) {  // 中文引号 " "
                return s.substring(1, s.length() - 1).strip();
            }
        }
        return s;
    }

    /**
     * 解析结果值对象。
     *
     * @param answer      LLM 回答正文（不含哨兵及之后的 suggestions）
     * @param suggestions 建议问题列表（0～3 条）
     */
    public record ParsedResponse(String answer, List<String> suggestions) {
    }
}