package com.easyagents.rag.retrieval;

import com.easyagents.core.document.Document;
import org.junit.Assert;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

/**
 * {@link RagScoreNormalizer} 回归测试。
 */
public class RagScoreNormalizerTest {

    /**
     * 验证关键词分数按有界函数归一化。
     */
    @Test
    public void shouldNormalizeKeywordScoresToZeroAndOneRange() {
        Document first = document(1, 9D, RagRetrievalMetadataKeys.KEYWORD_SCORE);
        Document second = document(2, 0D, RagRetrievalMetadataKeys.KEYWORD_SCORE);

        RagScoreNormalizer.normalize(Arrays.asList(first, second), RetrievalMode.KEYWORD, false);

        Assert.assertEquals(0.9D, first.getScore(), 0.0001D);
        Assert.assertEquals(0D, second.getScore(), 0.0001D);
    }

    /**
     * 验证混合检索 RRF 分数按理论上界归一化。
     */
    @Test
    public void shouldNormalizeHybridFusionScoreByRrfUpperBound() {
        Document document = document(1, 2D / (RrfFusionStrategy.DEFAULT_RRF_K + 1D), RagRetrievalMetadataKeys.FUSION_SCORE);

        RagScoreNormalizer.normalize(Arrays.asList(document), RetrievalMode.HYBRID, false);

        Assert.assertEquals(1D, document.getScore(), 0.0001D);
    }

    /**
     * 验证重排模型返回的绝对相关度保持不变。
     */
    @Test
    public void shouldPreserveRerankRelevanceScores() {
        List<Document> documents = Arrays.asList(
            document(1, 0.1D, RagRetrievalMetadataKeys.RERANK_SCORE),
            document(2, 0.5D, RagRetrievalMetadataKeys.RERANK_SCORE),
            document(3, 0.9D, RagRetrievalMetadataKeys.RERANK_SCORE)
        );

        RagScoreNormalizer.normalize(documents, RetrievalMode.HYBRID, true);

        Assert.assertEquals(0.1D, documents.get(0).getScore(), 0.0001D);
        Assert.assertEquals(0.5D, documents.get(1).getScore(), 0.0001D);
        Assert.assertEquals(0.9D, documents.get(2).getScore(), 0.0001D);
    }

    /**
     * 验证单条低分重排结果不会被抬高为高相关结果。
     */
    @Test
    public void shouldKeepSingleLowRerankScoreLow() {
        Document document = document(1, 0.1D, RagRetrievalMetadataKeys.RERANK_SCORE);

        RagScoreNormalizer.normalize(Arrays.asList(document), RetrievalMode.HYBRID, true);

        Assert.assertEquals(0.1D, document.getScore(), 0.0001D);
    }

    /**
     * 验证越界重排分数会被限制到零到一范围。
     */
    @Test
    public void shouldClampRerankScoresToZeroAndOneRange() {
        List<Document> documents = Arrays.asList(
            document(1, -0.2D, RagRetrievalMetadataKeys.RERANK_SCORE),
            document(2, 1.2D, RagRetrievalMetadataKeys.RERANK_SCORE)
        );

        RagScoreNormalizer.normalize(documents, RetrievalMode.HYBRID, true);

        Assert.assertEquals(0D, documents.get(0).getScore(), 0.0001D);
        Assert.assertEquals(1D, documents.get(1).getScore(), 0.0001D);
    }

    /**
     * 创建携带原始分数元数据的测试文档。
     *
     * @param id 文档 ID
     * @param score 原始分数
     * @param metadataKey 分数元数据键
     * @return 测试文档
     */
    private Document document(Object id, Double score, String metadataKey) {
        Document document = new Document();
        document.setId(id);
        document.setScore(score);
        document.addMetadata(metadataKey, score);
        return document;
    }
}
