// ScoringModel 核心方法
public interface ScoringModel {

    // 对单个文本段打分
    Response<Double> score(TextSegment segment, String query);

    // 批量打分（推荐使用，一次推理多条）
    Response<List<Double>> scoreAll(List<TextSegment> segments, String query);
}

// RerankingConfig.java
package com.example.rag.config;

import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.bgesmallenv15q.BgeSmallEnV15QuantizedEmbeddingModel;
import dev.langchain4j.model.scoring.ScoringModel;
import dev.langchain4j.model.scoring.onnx.OnnxScoringModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.content.aggregator.ContentAggregator;
import dev.langchain4j.rag.content.aggregator.ReRankingContentAggregator;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.data.segment.TextSegment;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class RerankingConfig {

    // 模型路径从 application.properties 注入
    @Value("${reranker.model.path}")
    private String modelPath;

    @Value("${reranker.tokenizer.path}")
    private String tokenizerPath;

    /**
     * 本地 ONNX 重排序模型 Bean
     * 使用 CPU 推理，无需 GPU，开箱即用
     */
    @Bean
    public ScoringModel scoringModel() {
        return new OnnxScoringModel(modelPath, tokenizerPath);
    }

    /**
     * 嵌入模型 Bean（用于向量化文档和查询）
     * bge-small-en-v1.5 量化版，24MB，纯本地推理
     */
    @Bean
    public EmbeddingModel embeddingModel() {
        return new BgeSmallEnV15QuantizedEmbeddingModel();
    }

    /**
     * 内存向量库（生产环境替换为 Milvus / Qdrant / PgVector 等）
     */
    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
        return new InMemoryEmbeddingStore<>();
    }

    /**
     * 向量检索器：粗排阶段，从库中捞出 Top-20 候选
     * maxResults(20) 故意设大，给重排序提供足够候选
     */
    @Bean
    public ContentRetriever contentRetriever(
            EmbeddingStore<TextSegment> embeddingStore,
            EmbeddingModel embeddingModel) {
        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(20)          // 粗排多捞，给精排留余量
                .minScore(0.5)           // 粗排阈值适当放宽
                .build();
    }

    /**
     * 重排序内容聚合器：精排阶段，从 Top-20 中取最优 Top-3
     *
     * minScore(0.7)：过滤掉精排后分数仍低的结果，防止低质量上下文污染 LLM
     * maxResults(3)：最终只保留 3 条最相关的内容送给 LLM
     */
    @Bean
    public ContentAggregator contentAggregator(ScoringModel scoringModel) {
        return ReRankingContentAggregator.builder()
                .scoringModel(scoringModel)
                .minScore(0.7)   // 精排分数阈值（0~1），低于此值的结果被丢弃
                .maxResults(3)   // 最终送给 LLM 的最大结果数
                .build();
    }

    /**
     * RAG 检索增强器：将粗排 + 精排串联为完整流水线
     */
    @Bean
    public RetrievalAugmentor retrievalAugmentor(
            ContentRetriever contentRetriever,
            ContentAggregator contentAggregator) {
        return DefaultRetrievalAugmentor.builder()
                .contentRetriever(contentRetriever)
                .contentAggregator(contentAggregator)
                .build();
    }
}

// StandaloneRerankingService.java
package com.example.rag.service;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.model.scoring.ScoringModel;
import org.springframework.stereotype.Service;

import java.util.Comparator;
import java.util.List;
import java.util.stream.IntStream;

@Service
public class StandaloneRerankingService {

    private final ScoringModel scoringModel;

    public StandaloneRerankingService(ScoringModel scoringModel) {
        this.scoringModel = scoringModel;
    }

    /**
     * 对候选文本列表按照与 query 的相关性重排序，返回 topK 条
     *
     * @param query      用户查询
     * @param candidates 粗排阶段返回的候选文本列表
     * @param topK       精排后保留的最大数量
     * @return 按相关性降序排列的 topK 结果
     */
    public List<ScoredText> rerank(String query, List<String> candidates, int topK) {
        // 将字符串包装为 TextSegment（LangChain4j 的文本段类型）
        List<TextSegment> segments = candidates.stream()
                .map(TextSegment::from)
                .toList();

        // 批量打分：一次推理为所有候选文本评分（比逐条调用效率高得多）
        Response<List<Double>> response = scoringModel.scoreAll(segments, query);
        List<Double> scores = response.content();

        // 将文本与分数绑定，按分数降序排列，取 topK
        return IntStream.range(0, candidates.size())
                .mapToObj(i -> new ScoredText(candidates.get(i), scores.get(i)))
                .sorted(Comparator.comparingDouble(ScoredText::score).reversed())
                .limit(topK)
                .toList();
    }

    /**
     * 打分结果封装（文本 + 相关性分数）
     */
    public record ScoredText(String text, double score) {}
}

// RagServiceFactory.java
package com.example.rag.service;

import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.service.AiServices;
import org.springframework.stereotype.Component;

@Component
public class RagServiceFactory {

    /**
     * 使用原始 AiServices API 手动构建 RAG 服务
     * 这里不用 langchain4j-spring 的 @AiService 注解，保持对流水线的完整控制
     */
    public Assistant createAssistant(
            ChatLanguageModel chatModel,
            RetrievalAugmentor retrievalAugmentor) {

        return AiServices.builder(Assistant.class)
                .chatLanguageModel(chatModel)
                .retrievalAugmentor(retrievalAugmentor)  // 注入带重排序的增强器
                .build();
    }

    // AI 服务接口（由 AiServices 动态代理实现）
    public interface Assistant {
        String chat(String userMessage);
    }
}

// RerankingTest.java
package com.example.rag;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.bgesmallenv15q.BgeSmallEnV15QuantizedEmbeddingModel;
import dev.langchain4j.model.scoring.ScoringModel;
import dev.langchain4j.model.scoring.onnx.OnnxScoringModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.content.aggregator.ContentAggregator;
import dev.langchain4j.rag.content.aggregator.DefaultContentAggregator;
import dev.langchain4j.rag.content.aggregator.ReRankingContentAggregator;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;

import java.util.Collection;
import java.util.List;
import java.util.Map;

import static org.assertj.core.api.Assertions.assertThat;

/**
 * 重排序功能集成测试
 *
 * 运行前提：
 * 1. 已从 HuggingFace 下载 BAAI/bge-reranker-base 的 ONNX 格式模型
 * 2. 设置环境变量或修改 MODEL_PATH / TOKENIZER_PATH 常量指向本地路径
 * 3. 如仅测试打分逻辑（不跑真实模型），可使用 MockScoringModel（见 6.2 节）
 */
class RerankingTest {

    // ⚠️ 修改为你的本地模型路径
    private static final String MODEL_PATH     = System.getenv().getOrDefault(
            "RERANKER_MODEL_PATH", "/models/bge-reranker-base/model.onnx");
    private static final String TOKENIZER_PATH = System.getenv().getOrDefault(
            "RERANKER_TOKENIZER_PATH", "/models/bge-reranker-base/tokenizer.json");

    private static EmbeddingModel      embeddingModel;
    private static EmbeddingStore<TextSegment> embeddingStore;
    private static ContentRetriever    contentRetriever;

    @BeforeAll
    static void setUp() {
        // 初始化嵌入模型（本地 ONNX，无需外部 API）
        embeddingModel = new BgeSmallEnV15QuantizedEmbeddingModel();
        embeddingStore = new InMemoryEmbeddingStore<>();

        // 准备知识库：写入几条故意设计的"相似但相关性不同"的文本
        List<Document> docs = List.of(
            Document.from("Java 线程池的 shutdown() 方法会平滑地关闭线程池，等待已提交的任务全部完成后再终止。"),
            Document.from("Python 的 ThreadPoolExecutor 可以通过 shutdown(wait=True) 关闭线程池。"),
            Document.from("Java 线程池 shutdownNow() 方法会立即尝试停止所有活跃任务，并返回等待中的任务列表。"),
            Document.from("Java 线程池最佳实践：使用 try-finally 确保 shutdown() 总是被调用，避免资源泄漏。"),
            Document.from("线程池的核心参数：corePoolSize、maximumPoolSize、keepAliveTime 和 BlockingQueue。"),
            Document.from("Redis 缓存穿透可以通过布隆过滤器（Bloom Filter）在数据库查询前过滤掉不存在的 key。")
        );

        // 使用 EmbeddingStoreIngestor 批量向量化并存入内存库
        DocumentSplitter splitter = DocumentSplitters.recursive(200, 20);
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(splitter)
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .build();
        docs.forEach(ingestor::ingest);

        // 粗排检索器：最多返回 10 条候选
        contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(10)
                .build();
    }

    @Test
    @DisplayName("对比：有无重排序的检索结果差异（通过直接调用 ContentRetriever + ScoringModel 验证）")
    void testRerankingVsNoReranking() {
        String query = "Java 如何安全地关闭线程池？";

        // ─── 场景 A：无重排序，直接取向量检索 Top-3 ───
        List<dev.langchain4j.rag.content.Content> rawResults =
                contentRetriever.retrieve(Query.from(query));

        System.out.println("=== 无重排序结果（向量相似度排序，取前 3）===");
        rawResults.stream().limit(3).forEach(c ->
                System.out.println("  · " + c.textSegment().text()));

        // ─── 场景 B：有重排序，用 ScoringModel 对向量检索结果精排 ───
        ScoringModel scoringModel = new OnnxScoringModel(MODEL_PATH, TOKENIZER_PATH);

        // 提取文本段列表
        List<TextSegment> candidates = rawResults.stream()
                .map(dev.langchain4j.rag.content.Content::textSegment)
                .toList();

        // 批量打分
        List<Double> scores = scoringModel.scoreAll(candidates, query).content();

        // 按分数降序排列，过滤 minScore < 0.6 的结果，取 Top-3
        List<String> rerankResults = IntStream.range(0, candidates.size())
                .boxed()
                .filter(i -> scores.get(i) >= 0.6)          // minScore 过滤
                .sorted((a, b) -> Double.compare(scores.get(b), scores.get(a)))
                .limit(3)
                .map(i -> String.format("[%.4f] %s", scores.get(i), candidates.get(i).text()))
                .toList();

        System.out.println("=== 有重排序结果（交叉编码器精排，minScore=0.6）===");
        rerankResults.forEach(r -> System.out.println("  · " + r));

        // 验证：重排序后结果数量不超过 3
        assertThat(rerankResults.size()).isLessThanOrEqualTo(3);

        // 验证：重排序后第一条结果的分数应高于 0.6
        if (!rerankResults.isEmpty()) {
            assertThat(rerankResults.get(0)).contains("Java");
        }
    }

    @Test
    @DisplayName("直接使用 ScoringModel 打分验证")
    void testScoringModelDirectly() {
        ScoringModel scoringModel = new OnnxScoringModel(MODEL_PATH, TOKENIZER_PATH);

        String query = "Java 如何安全地关闭线程池？";
        List<TextSegment> candidates = List.of(
            TextSegment.from("Java 线程池的 shutdown() 方法会平滑地关闭线程池。"),     // 高度相关
            TextSegment.from("Python 的 ThreadPoolExecutor 可以通过 shutdown() 关闭。"), // 语言不符
            TextSegment.from("Redis 缓存穿透的解决方案是使用布隆过滤器。")                // 完全不相关
        );

        // 批量打分
        List<Double> scores = scoringModel.scoreAll(candidates, query).content();

        System.out.println("=== 相关性打分结果 ===");
        for (int i = 0; i < candidates.size(); i++) {
            System.out.printf("  分数: %.4f | %s%n",
                    scores.get(i), candidates.get(i).text());
        }

        // 验证：Java 线程池关闭的文本分数应明显高于 Python 和 Redis
        double javaScore  = scores.get(0);
        double pythonScore = scores.get(1);
        double redisScore  = scores.get(2);

        assertThat(javaScore).isGreaterThan(pythonScore);
        assertThat(javaScore).isGreaterThan(redisScore);
    }
}

// MockScoringModelTest.java
package com.example.rag;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.model.scoring.ScoringModel;
import dev.langchain4j.rag.content.aggregator.ContentAggregator;
import dev.langchain4j.rag.content.aggregator.ReRankingContentAggregator;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;

import java.util.List;
import java.util.Map;

import static org.assertj.core.api.Assertions.assertThat;

class MockScoringModelTest {

    @Test
    @DisplayName("使用 Mock ScoringModel 验证重排序聚合逻辑")
    void testWithMockScoringModel() {
        // Mock ScoringModel：根据文本内容返回预设分数（不调用真实模型）
        Map<String, Double> mockScores = Map.of(
            "Java shutdown() 关闭线程池", 0.95,
            "Python 关闭线程池",          0.40,
            "Redis 缓存穿透解决方案",      0.10
        );

        ScoringModel mockScoringModel = new ScoringModel() {
            @Override
            public Response<Double> score(TextSegment segment, String query) {
                double score = mockScores.getOrDefault(segment.text(), 0.5);
                return Response.from(score);
            }

            @Override
            public Response<List<Double>> scoreAll(List<TextSegment> segments, String query) {
                List<Double> scores = segments.stream()
                        .map(s -> mockScores.getOrDefault(s.text(), 0.5))
                        .toList();
                return Response.from(scores);
            }
        };

        // 构建 ReRankingContentAggregator，minScore=0.6 会过滤掉 Python 和 Redis 相关内容
        ContentAggregator aggregator = ReRankingContentAggregator.builder()
                .scoringModel(mockScoringModel)
                .minScore(0.6)
                .maxResults(2)
                .build();

        // 验证 aggregator 已正确创建（实际内容聚合需要配合 RetrievalAugmentor 使用）
        assertThat(aggregator).isNotNull();

        // 直接测试打分逻辑
        List<TextSegment> candidates = List.of(
            TextSegment.from("Java shutdown() 关闭线程池"),
            TextSegment.from("Python 关闭线程池"),
            TextSegment.from("Redis 缓存穿透解决方案")
        );

        List<Double> scores = mockScoringModel.scoreAll(candidates, "Java 关闭线程池").content();

        // 验证分数排序符合预期
        assertThat(scores.get(0)).isGreaterThan(scores.get(1)); // Java > Python
        assertThat(scores.get(0)).isGreaterThan(scores.get(2)); // Java > Redis
        assertThat(scores.get(0)).isGreaterThanOrEqualTo(0.6);  // Java 通过阈值
        assertThat(scores.get(1)).isLessThan(0.6);              // Python 被过滤
    }
}

// ❌ 错误做法：粗排和精排的 maxResults 设置相同
ContentRetriever retriever = EmbeddingStoreContentRetriever.builder()
        .maxResults(3)    // 粗排只捞 3 条
        .build();
ContentAggregator aggregator = ReRankingContentAggregator.builder()
        .maxResults(3)    // 精排也要 3 条，等于没有过滤
        .build();

// ✅ 正确做法：粗排多捞，精排严格过滤
ContentRetriever retriever = EmbeddingStoreContentRetriever.builder()
        .maxResults(20)   // 粗排宽松召回，给精排足够候选
        .minScore(0.4)    // 粗排阈值适当放宽，避免错过相关结果
        .build();
ContentAggregator aggregator = ReRankingContentAggregator.builder()
        .scoringModel(scoringModel)
        .maxResults(3)    // 精排严格，只取最优 3 条送 LLM
        .minScore(0.7)    // 精排阈值严格，过滤低质量上下文
        .build();

// 从本地 ONNX 切换到 Cohere 云端 API，只需换这一个 Bean
// （需要添加 langchain4j-cohere 依赖）
@Bean
public ScoringModel scoringModel() {
    // 本地模式：return new OnnxScoringModel(modelPath, tokenizerPath);
    return CohereScoringModel.builder()
            .apiKey(System.getenv("COHERE_API_KEY"))
            .modelName("rerank-multilingual-v3.0")
            .build();
}