RAG/src/main/java/com/wok/supportbot/rag/RagPipeline.java

package com.wok.supportbot.rag;

import com.wok.supportbot.app.ChatContext;
import com.wok.supportbot.chatmemory.DatabaseChatMemory;
import com.wok.supportbot.config.RagPromptConfig;
import com.wok.supportbot.rag.preretrieval.CompressionQueryRewriter;
import com.wok.supportbot.rag.preretrieval.MultiQueryExpanderRewriter;
import com.wok.supportbot.rag.preretrieval.RewriteQueryRewriter;
import com.wok.supportbot.rag.preretrieval.TranslationQueryRewriter;
import com.wok.supportbot.service.FaqMatchEngine;
import com.wok.supportbot.service.FaqMatchEngine.FaqMatchResult;
import com.wok.supportbot.service.RagHitLogService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Collectors;

/**
 * 统一 RAG 检索管道。
 * <p>
 * 收敛原本分散在 {@code AssistantApp} 中按策略分支的检索逻辑：
 * <ul>
 *   <li>FAQ 优先匹配（复用 {@link FaqMatchEngine} 三级匹配）</li>
 *   <li>查询重写（按 {@code rewriteStrategy} 复用 {@code rag/preretrieval/*} 四种 rewriter）</li>
 *   <li>统一检索：{@code MULTI_QUERY} 扩展多查询后按文档 ID 去重合并，其余策略单查询检索</li>
 *   <li>统一资料块模板 {@link #buildRagContextBlock}，替代原 {@code buildRetrievalAdvisor.qaTemplate}
 *       与 {@code buildRagSystemPrompt} 两份回答模板</li>
 * </ul>
 * <p>
 * 统一为"手动检索 + 资料块注入系统提示词"模式：所有策略都先由本管道检索出文档，
 * 再由调用方把 {@link RagContext#contextText()} 组装进系统提示词，
 * 不再使用 {@code RetrievalAugmentationAdvisor} 的 query augmenter 自动注入，
 * 消除上下文注入位置随策略不同而不同的不一致。
 * <p>
 * 阶段一作为旁路组件存在，旧 {@code AssistantApp} RAG 路径未改动；阶段二由 {@code ChatPipeline} 接入。
 * <p>
 * {@code @pipeline} rag-layer order=1<br>
 * {@code @pipeline-step} retrieve: FAQ优先匹配 → 查询重写/扩展 → similaritySearch(PGVector) → 资料拼接<br>
 * {@code @pipeline-step} similaritySearch: 纯向量检索 topK=4 + CategoryFilter 分类过滤<br>
 * 注意: HybridSearchService/RrfFusion/RerankerService 尚未接入本管道，当前仅单路向量检索。<br>
 * 同步至: frontend/src/views/PipelineFlow.vue RAG 子图
 */
@Component
@Slf4j
public class RagPipeline {

    /** 单路检索 topK，与原 AssistantApp 保持一致 */
    private static final int TOP_K = 4;

    /** 多路检索合并后的文档封顶数，避免上下文膨胀 */
    private static final int MAX_DOCS = 8;

    /** COMPRESSION 策略读取的对话历史条数 */
    private static final int COMPRESSION_HISTORY_SIZE = 10;

    @Resource
    private VectorStore pgVectorVectorStore;

    @Resource
    private FaqMatchEngine faqMatchEngine;

    @Resource
    private RagPromptConfig ragPromptConfig;

    @Resource
    private RagHitLogService ragHitLogService;

    @Resource
    private CategoryFilter categoryFilter;

    @Resource
    private RewriteQueryRewriter rewriteQueryRewriter;

    @Resource
    private TranslationQueryRewriter translationQueryRewriter;

    @Resource
    private CompressionQueryRewriter compressionQueryRewriter;

    @Resource
    private MultiQueryExpanderRewriter multiQueryExpanderRewriter;

    /** MULTI_QUERY 多路检索扇出线程池：并行执行各路向量检索，饱和时退回请求线程串行（CallerRuns） */
    @Resource(name = "ragRetrieveExecutor")
    private ThreadPoolTaskExecutor ragRetrieveExecutor;

    /** MULTI_QUERY 多路检索是否并行执行（false 回退串行，供高并发/embedding 侧限流时降级） */
    @Value("${knowledge.rag.multiquery.parallel:true}")
    private boolean multiQueryParallel;

    private final DatabaseChatMemory chatMemory;

    public RagPipeline(DatabaseChatMemory chatMemory) {
        this.chatMemory = chatMemory;
    }

    /**
     * 执行一次统一的 RAG 检索。
     * <p>
     * 流程：FAQ 优先 → 查询重写/扩展 → 统一检索 → 拼接资料文本。
     *
     * @param ctx 对话上下文（使用 {@code message / chatId / rewriteStrategy / categoryIds}）
     * @return 检索结果；FAQ 命中时 documents 与 contextText 为空，rewrittenQuery 为原始 message
     */
    public RagContext retrieve(ChatContext ctx) {
        return retrieve(ctx, false);
    }

    /**
     * 执行一次统一的 RAG 检索（可跳过前序已做过的 FAQ 匹配）。
     * <p>
     * 流程：FAQ 优先（未匹配过时）→ 查询重写/扩展 → 统一检索 → 拼接资料文本。
     *
     * @param ctx               对话上下文（使用 {@code message / chatId / rewriteStrategy / categoryIds}）
     * @param faqAlreadyMatched 编排层是否已在前序阶段（FAQ 高置信未命中降级）干净跑过完整 FAQ 三级匹配；
     *                          true 时跳过 retrieve 内重复的 FAQ 匹配，避免同一请求重复做 FAQ 语义 embedding
     * @return 检索结果；FAQ 命中时 documents 与 contextText 为空，rewrittenQuery 为原始 message
     */
    public RagContext retrieve(ChatContext ctx, boolean faqAlreadyMatched) {
        // 1. FAQ 优先匹配：命中则直接返回标准答案，跳过检索与生成（已匹配过则跳过，防止重复 embedding）
        if (!faqAlreadyMatched) {
            Optional<FaqMatchResult> faqMatch = tryFaqMatchResult(ctx.message(), ctx.categoryIds());
            if (faqMatch.isPresent()) {
                log.info("FAQ 命中，跳过知识库检索: chatId={}, matchType={}", ctx.chatId(), faqMatch.get().getMatchType());
                String answer = faqMatch.get().getFaq().getAnswer();
                return new RagContext(Optional.ofNullable(answer), Collections.emptyList(), "", ctx.message(),
                        currentSearchMode(), faqMatch.get());
            }
        }

        // 2. 统一检索 + 组装结果
        return retrieveDocumentsAndContext(ctx);
    }

    /**
     * 仅检索知识库片段（跳过 FAQ 匹配），用于"引用来源"展示。
     * <p>
     * 与 {@link #retrieve} 共用同一套查询重写与检索逻辑，确保来源即答案所依据的片段，
     * 但不触发 FAQ 优先匹配——来源接口的语义是展示 KB 片段，FAQ 命中时本就无 KB 来源。
     *
     * @param ctx 对话上下文
     * @return 命中的知识库片段（含 metadata），无命中返回空列表
     */
    public List<Document> retrieveDocuments(ChatContext ctx) {
        return retrieveDocumentsAndContext(ctx).documents();
    }

    /**
     * 统一检索 + 拼接资料文本（不含 FAQ 匹配），供 {@link #retrieve} 与 {@link #retrieveDocuments} 复用。
     */
    private RagContext retrieveDocumentsAndContext(ChatContext ctx) {
        List<Document> docs;
        String rewrittenQuery;
        if ("MULTI_QUERY".equalsIgnoreCase(ctx.rewriteStrategy())) {
            // MULTI_QUERY：资料注入 system，user 消息保持原始 message
            docs = retrieveMultiQueryDocs(ctx.message(), ctx.categoryIds());
            rewrittenQuery = ctx.message();
        } else {
            rewrittenQuery = rewriteQuery(ctx.message(), ctx.chatId(), ctx.rewriteStrategy());
            docs = similaritySearch(rewrittenQuery, ctx.categoryIds());
        }

        // RAG 命中日志：记录本次检索的命中/未命中情况
        logRagHit(ctx.chatId(), ctx.message(), docs, currentSearchMode());

        String contextText = joinContext(docs);
        return new RagContext(Optional.empty(), docs, contextText, rewrittenQuery, currentSearchMode(), null);
    }

    /**
     * 构造统一的 RAG 资料块文本，供调用方追加到系统提示词。
     * <p>
     * 替代原 {@code buildRetrievalAdvisor} 的 qaTemplate 与 {@code buildRagSystemPrompt} 两份模板，
     * 所有策略共用同一份"回答硬性规则 + 知识库资料"格式。
     *
     * @param contextText 检索拼接的资料文本，为空时返回空串（调用方据此决定是否追加）
     */
    public String buildRagContextBlock(String contextText) {
        if (!StringUtils.hasText(contextText)) {
            return "";
        }
        return "\n\n【RAG回答硬性规则】\n" + ragPromptConfig.getAnswerRules()
                + "\n\n【知识库资料】\n" + contextText;
    }

    // ==================== FAQ 匹配 ====================

    /**
     * 尝试 FAQ 三级匹配（精确→关键词→语义），命中返回完整匹配结果（含 matchType/score）。
     * 异常时降级为未命中。
     *
     * @param message     用户问题
     * @param categoryIds 角色授权分类 ID 列表（null/空表示不限制）
     */
    public Optional<FaqMatchResult> tryFaqMatchResult(String message, List<Long> categoryIds) {
        try {
            return faqMatchEngine.match(message, categoryIds);
        } catch (Exception e) {
            log.warn("FAQ 匹配异常，降级到 RAG: {}", e.getMessage());
            return Optional.empty();
        }
    }

    /**
     * FAQ 匹配结果 + 是否「完整跑完三级匹配后仍未命中」（区别于异常降级）。
     * 供编排层判定：前序已干净跑过 FAQ 三级匹配且未命中时，可在后续检索中跳过重复的 FAQ 匹配。
     */
    public record FaqMatchOutcome(Optional<FaqMatchResult> result, boolean completedCleanly) {
    }

    /**
     * 尝试 FAQ 三级匹配并携带「是否干净完成」标记。
     * match 内部各子级已吞掉各自异常并返回 empty，此处仅在 match 整体抛出未捕获异常时置 completedCleanly=false，
     * 避免异常降级被误当作「已匹配未命中」而在后续检索中跳过第二次 FAQ（可能本可命中）。
     */
    public FaqMatchOutcome tryFaqMatchClean(String message, List<Long> categoryIds) {
        try {
            return new FaqMatchOutcome(faqMatchEngine.match(message, categoryIds), true);
        } catch (Exception e) {
            log.warn("FAQ 匹配异常(clean 标记)，降级到 RAG: {}", e.getMessage());
            return new FaqMatchOutcome(Optional.empty(), false);
        }
    }

    /**
     * 尝试 FAQ 三级匹配，命中返回标准答案（仅答案文本）。
     * 异常时降级为未命中，供仅需答案的调用方使用。
     *
     * @param message     用户问题
     * @param categoryIds 角色授权分类 ID 列表（null/空表示不限制）
     */
    public Optional<String> tryFaqMatch(String message, List<Long> categoryIds) {
        return tryFaqMatchResult(message, categoryIds).map(result -> result.getFaq().getAnswer());
    }

    /**
     * 当前主管道的真实检索模式。
     * <p>
     * 目前仅单路向量检索（VECTOR）；HybridSearchService 接入主管道后，
     * 此方法应改为根据上下文透传 VECTOR / KEYWORD / HYBRID。
     */
    private String currentSearchMode() {
        return "VECTOR";
    }

    // ==================== 查询重写 ====================

    /**
     * 按策略改写查询（MULTI_QUERY 不走此方法，由 {@link #retrieveMultiQueryDocs} 自行扩展）。
     * 重写失败时降级为原始查询，避免 RAG_REWRITE 配置异常导致整次对话不可用。
     */
    private String rewriteQuery(String message, String chatId, String strategy) {
        if (strategy == null || strategy.isEmpty()) {
            return message;
        }
        try {
            switch (strategy.toUpperCase()) {
                case "REWRITE":
                    return rewriteQueryRewriter.doQueryRewrite(message);
                case "TRANSLATION":
                    return translationQueryRewriter.doQueryRewrite(message);
                case "COMPRESSION":
                    // 查询压缩需要对话历史，从会话记忆取最近若干条，把指代不清的追问补全为独立查询
                    List<Message> history = chatMemory.get(chatId, COMPRESSION_HISTORY_SIZE);
                    return compressionQueryRewriter.doQueryRewrite(message, history);
                case "NONE":
                default:
                    return message;
            }
        } catch (Exception e) {
            log.warn("查询重写失败 [strategy={}, chatId={}]，降级使用原始查询: {}", strategy, chatId, e.getMessage());
            return message;
        }
    }

    // ==================== 多路检索 ====================

    /**
     * 多路查询扩展 + 分别检索 + 按文档 ID 去重合并。
     * 扩展失败时退回原问题检索，避免整次 RAG 不可用。
     * 多路检索默认并行执行（每路 embedding HTTP + PG 查询相互独立），关闭或异常时回退串行。
     */
    private List<Document> retrieveMultiQueryDocs(String message, List<Long> categoryIds) {
        List<String> expandedQueries;
        try {
            expandedQueries = multiQueryExpanderRewriter.doQueryRewrite(message);
        } catch (Exception e) {
            log.warn("多路查询扩展失败，降级为原始问题检索: {}", e.getMessage());
            expandedQueries = List.of(message);
        }
        if (expandedQueries == null || expandedQueries.isEmpty()) {
            expandedQueries = List.of(message);
        }
        log.info("多路查询扩展结果: {}", expandedQueries);

        // 多路检索默认并行执行（每路 embedding HTTP + PG 查询相互独立，可把 N 路串行降为 1 段最慢路）；
        // 单路或关闭并行开关时回退串行，保证结果与旧版一致
        boolean useParallel = multiQueryParallel && expandedQueries.size() > 1;
        return useParallel
                ? retrieveMultiQueryParallel(expandedQueries, categoryIds)
                : retrieveMultiQuerySequential(expandedQueries, categoryIds);
    }

    /**
     * 串行执行多路检索（并行开关关闭 / 单路 / 异常回退路径，行为与旧版一致）。
     */
    private List<Document> retrieveMultiQuerySequential(List<String> expandedQueries, List<Long> categoryIds) {
        Map<String, Document> merged = new LinkedHashMap<>();
        for (String query : expandedQueries) {
            if (!StringUtils.hasText(query) || merged.size() >= MAX_DOCS) {
                continue;
            }
            List<Document> docs = similaritySearch(query, categoryIds);
            for (Document doc : docs) {
                if (merged.size() >= MAX_DOCS) {
                    break;
                }
                merged.putIfAbsent(doc.getId(), doc);
            }
        }
        return new ArrayList<>(merged.values());
    }

    /**
     * 并行执行多路检索：各路相互独立（embedding HTTP + PG 查询），可把 N 路串行降为 1 段最慢路。
     * join 后仍按原 query 顺序合并去重（LinkedHashMap + MAX_DOCS 封顶），保证文档集合与串行结果一致；
     * 任一路异常降级为空列表，不拖垮整次检索。
     */
    private List<Document> retrieveMultiQueryParallel(List<String> expandedQueries, List<Long> categoryIds) {
        List<CompletableFuture<List<Document>>> futures = new ArrayList<>(expandedQueries.size());
        for (String query : expandedQueries) {
            if (!StringUtils.hasText(query)) {
                futures.add(CompletableFuture.completedFuture(Collections.emptyList()));
                continue;
            }
            futures.add(CompletableFuture
                    .supplyAsync(() -> similaritySearch(query, categoryIds), ragRetrieveExecutor)
                    .exceptionally(ex -> {
                        log.warn("多路并行检索单路失败，该路降级为空: error={}", ex.getMessage());
                        return Collections.emptyList();
                    }));
        }
        CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();

        // 按原 query 顺序合并，去重/封顶规则与串行路径一致（不按完成先后，保证结果确定）
        Map<String, Document> merged = new LinkedHashMap<>();
        for (int i = 0; i < expandedQueries.size() && merged.size() < MAX_DOCS; i++) {
            if (!StringUtils.hasText(expandedQueries.get(i))) {
                continue;
            }
            for (Document doc : futures.get(i).join()) {
                if (merged.size() >= MAX_DOCS) {
                    break;
                }
                merged.putIfAbsent(doc.getId(), doc);
            }
        }
        return new ArrayList<>(merged.values());
    }

    // ==================== 底层检索 ====================

    /**
     * 单查询向量检索，附带分类过滤。
     */
    private List<Document> similaritySearch(String query, List<Long> categoryIds) {
        if (!StringUtils.hasText(query)) {
            return Collections.emptyList();
        }
        SearchRequest.Builder builder = SearchRequest.builder()
                .similarityThreshold(0.0)
                .topK(TOP_K);
        Filter.Expression filterExpression = categoryFilter.buildExpression(categoryIds);
        if (filterExpression != null) {
            builder.filterExpression(filterExpression);
        }
        List<Document> docs = pgVectorVectorStore.similaritySearch(builder.query(query).build());
        return docs != null ? docs : Collections.emptyList();
    }

    /**
     * 把文档列表拼接为资料文本，过滤空内容，用分隔符区分片段。
     */
    private String joinContext(List<Document> docs) {
        if (docs == null || docs.isEmpty()) {
            return "";
        }
        return docs.stream()
                .map(Document::getText)
                .filter(StringUtils::hasText)
                .collect(Collectors.joining("\n\n---\n\n"));
    }

    /**
     * 异步记录 RAG 命中日志到 rag_hit_log 表。
     * 有命中文档时逐条记录 hit，无命中时记录一条 miss。
     */
    private void logRagHit(String conversationId, String userQuery, List<Document> docs, String searchMode) {
        try {
            if (docs == null || docs.isEmpty()) {
                ragHitLogService.recordMiss(conversationId, userQuery, searchMode);
            } else {
                for (Document doc : docs) {
                    String docIdStr = String.valueOf(doc.getMetadata().getOrDefault("documentId", ""));
                    Long documentId = null;
                    try {
                        if (!docIdStr.isEmpty()) documentId = Long.parseLong(docIdStr);
                    } catch (NumberFormatException ignored) { }
                    String title = String.valueOf(doc.getMetadata().getOrDefault("title", ""));
                    String score = String.valueOf(doc.getMetadata().getOrDefault("score", ""));
                    ragHitLogService.recordHit(conversationId, userQuery, documentId, title, score, searchMode);
                }
            }
        } catch (Exception e) {
            log.warn("RAG 命中日志记录失败: {}", e.getMessage());
        }
    }
}