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 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.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.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} 接入。
 */
@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 CategoryFilter categoryFilter;

    @Resource
    private RewriteQueryRewriter rewriteQueryRewriter;

    @Resource
    private TranslationQueryRewriter translationQueryRewriter;

    @Resource
    private CompressionQueryRewriter compressionQueryRewriter;

    @Resource
    private MultiQueryExpanderRewriter multiQueryExpanderRewriter;

    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) {
        // 1. FAQ 优先匹配：命中则直接返回标准答案，跳过检索与生成
        Optional<String> faqAnswer = tryFaqMatch(ctx.message());
        if (faqAnswer.isPresent()) {
            log.info("FAQ 命中，跳过知识库检索: chatId={}", ctx.chatId());
            return new RagContext(faqAnswer, Collections.emptyList(), "", ctx.message());
        }

        // 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());
        }
        String contextText = joinContext(docs);
        return new RagContext(Optional.empty(), docs, contextText, rewrittenQuery);
    }

    /**
     * 构造统一的 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 三级匹配（精确→关键词→语义），命中返回标准答案。
     * 异常时降级为未命中，继续走 RAG 检索。
     */
    private Optional<String> tryFaqMatch(String message) {
        try {
            return faqMatchEngine.match(message)
                    .map(result -> result.getFaq().getAnswer());
        } catch (Exception e) {
            log.warn("FAQ 匹配异常，降级到 RAG: {}", e.getMessage());
            return Optional.empty();
        }
    }

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

    /**
     * 按策略改写查询（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 不可用。
     */
    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);

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

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

    /**
     * 单查询向量检索，附带分类过滤。
     */
    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"));
    }
}