package com.easyagents.rag.retrieval;

import com.easyagents.core.document.Document;

import java.util.ArrayList;
import java.util.List;

public final class RagScoreNormalizer {

    private RagScoreNormalizer() {
    }

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

    private static void normalizeRerankScores(List<Document> documents) {
        List<Double> rawScores = new ArrayList<Double>(documents.size());
        boolean allPresent = true;
        Double min = null;
        Double max = null;
        for (Document document : documents) {
            Double rawScore = readRawScore(document, RagRetrievalMetadataKeys.RERANK_SCORE, document == null ? null : document.getScore());
            rawScores.add(rawScore);
            if (rawScore == null) {
                allPresent = false;
                continue;
            }
            min = min == null ? rawScore : Math.min(min, rawScore);
            max = max == null ? rawScore : Math.max(max, rawScore);
        }

        if (allPresent && min != null && max != null && Double.compare(max, min) != 0) {
            for (int i = 0; i < documents.size(); i++) {
                Double rawScore = rawScores.get(i);
                documents.get(i).setScore(clamp01((rawScore - min) / (max - min)));
            }
            return;
        }

        if (documents.size() == 1) {
            documents.get(0).setScore(1D);
            return;
        }

        int size = documents.size();
        for (int i = 0; i < size; i++) {
            documents.get(i).setScore(clamp01(1D - ((double) i / (double) (size - 1))));
        }
    }

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

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