package com.easyagents.rag.retrieval;

import com.easyagents.core.document.Document;

import java.util.List;

/**
 * 将不同检索路径的原始分数转换为统一的零到一最终相关度。
 */
public final class RagScoreNormalizer {

    /**
     * 禁止实例化工具类。
     */
    private RagScoreNormalizer() {
    }

    /**
     * 按检索模式归一化文档最终分数。
     *
     * @param documents 待归一化文档
     * @param retrievalMode 检索模式
     * @param reranked 是否已经过重排模型
     */
    public static void normalize(List<Document> documents, RetrievalMode retrievalMode, boolean reranked) {
        if (documents == null || documents.isEmpty()) {
            return;
        }
        if (reranked) {
            normalizeRerankScores(documents);
            return;
        }

        RetrievalMode effectiveMode = retrievalMode == null ? RetrievalMode.HYBRID : retrievalMode;
        if (effectiveMode == RetrievalMode.VECTOR) {
            for (Document document : documents) {
                document.setScore(clamp01(readRawScore(document, RagRetrievalMetadataKeys.VECTOR_SCORE, document == null ? null : document.getScore())));
            }
            return;
        }
        if (effectiveMode == RetrievalMode.KEYWORD) {
            for (Document document : documents) {
                Double rawScore = readRawScore(document, RagRetrievalMetadataKeys.KEYWORD_SCORE, document == null ? null : document.getScore());
                if (rawScore == null || rawScore <= 0D) {
                    document.setScore(0D);
                } else {
                    document.setScore(clamp01(rawScore / (1D + rawScore)));
                }
            }
            return;
        }

        double maxHybridRrfScore = 2D / (RrfFusionStrategy.DEFAULT_RRF_K + 1D);
        for (Document document : documents) {
            Double rawScore = readRawScore(document, RagRetrievalMetadataKeys.FUSION_SCORE, document == null ? null : document.getScore());
            if (rawScore == null || rawScore <= 0D) {
                document.setScore(0D);
            } else {
                document.setScore(clamp01(rawScore / maxHybridRrfScore));
            }
        }
    }

    /**
     * 保留重排模型返回的绝对相关度，并将异常范围限制到零到一。
     *
     * @param documents 待归一化的文档
     */
    private static void normalizeRerankScores(List<Document> documents) {
        for (Document document : documents) {
            Double rawScore = readRawScore(document, RagRetrievalMetadataKeys.RERANK_SCORE, document == null ? null : document.getScore());
            if (document != null) {
                // Rerank 适配器返回绝对相关度，按查询结果集再次缩放会把低相关第一名错误抬高到 1。
                document.setScore(clamp01(rawScore));
            }
        }
    }

    /**
     * 优先读取指定元数据中的原始分数。
     *
     * @param document 文档
     * @param metadataKey 分数元数据键
     * @param fallback 元数据不可用时的回退分数
     * @return 原始分数；文档为空时返回 null
     */
    private static Double readRawScore(Document document, String metadataKey, Double fallback) {
        if (document == null) {
            return null;
        }
        Object metadataValue = document.getMetadata(metadataKey);
        if (metadataValue instanceof Number) {
            return ((Number) metadataValue).doubleValue();
        }
        if (metadataValue instanceof String) {
            try {
                return Double.valueOf((String) metadataValue);
            } catch (NumberFormatException ignore) {
                return fallback;
            }
        }
        return fallback;
    }

    /**
     * 将可空分数限制到零到一范围。
     *
     * @param value 原始分数
     * @return 有效最终分数
     */
    private static double clamp01(Double value) {
        if (value == null || value.isNaN() || value.isInfinite()) {
            return 0D;
        }
        if (value < 0D) {
            return 0D;
        }
        if (value > 1D) {
            return 1D;
        }
        return value;
    }
}
