package com.easyagents.rag.retrieval;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

public class RrfFusionStrategy implements FusionStrategy {

    public static final int DEFAULT_RRF_K = 60;

    private final int rrfK;

    public RrfFusionStrategy() {
        this(DEFAULT_RRF_K);
    }

    public RrfFusionStrategy(int rrfK) {
        this.rrfK = rrfK <= 0 ? DEFAULT_RRF_K : rrfK;
    }

    @Override
    public List<RagHit> fuse(List<RagHit> vectorHits, List<RagHit> keywordHits, int topK) {
        Map<String, RagHit> mergedHits = new LinkedHashMap<String, RagHit>();
        Map<String, Double> fusionScores = new LinkedHashMap<String, Double>();

        mergeHits(mergedHits, fusionScores, vectorHits, HitSource.VECTOR);
        mergeHits(mergedHits, fusionScores, keywordHits, HitSource.KEYWORD);

        List<RagHit> results = new ArrayList<RagHit>(mergedHits.values());
        Collections.sort(results, new Comparator<RagHit>() {
            @Override
            public int compare(RagHit left, RagHit right) {
                int scoreCompare = Double.compare(nullSafeScore(right.getScore()), nullSafeScore(left.getScore()));
                if (scoreCompare != 0) {
                    return scoreCompare;
                }
                return String.valueOf(left.getDocumentId()).compareTo(String.valueOf(right.getDocumentId()));
            }
        });
        assignRank(results);
        if (topK > 0 && results.size() > topK) {
            return new ArrayList<RagHit>(results.subList(0, topK));
        }
        return results;
    }

    private void mergeHits(Map<String, RagHit> mergedHits,
                           Map<String, Double> fusionScores,
                           List<RagHit> incomingHits,
                           HitSource source) {
        if (incomingHits == null) {
            return;
        }
        for (int index = 0; index < incomingHits.size(); index++) {
            RagHit incoming = incomingHits.get(index);
            if (incoming == null || incoming.getDocumentId() == null) {
                continue;
            }
            String key = String.valueOf(incoming.getDocumentId());
            RagHit merged = mergedHits.get(key);
            if (merged == null) {
                merged = incoming.copy();
                merged.setHitSource(source);
                mergedHits.put(key, merged);
                fusionScores.put(key, 0D);
            } else {
                mergeBaseFields(merged, incoming);
                merged.setHitSource(HitSource.merge(merged.getHitSource(), source));
            }

            int rank = index + 1;
            if (source == HitSource.VECTOR) {
                merged.setVectorScore(incoming.getVectorScore() != null ? incoming.getVectorScore() : incoming.getScore());
                merged.getMetadata().put(RagRetrievalMetadataKeys.VECTOR_RANK, rank);
            } else {
                merged.setKeywordScore(incoming.getKeywordScore() != null ? incoming.getKeywordScore() : incoming.getScore());
                merged.getMetadata().put(RagRetrievalMetadataKeys.KEYWORD_RANK, rank);
            }
            double fusedScore = fusionScores.get(key) + (1D / (rrfK + rank));
            fusionScores.put(key, fusedScore);
            merged.setScore(fusedScore);
        }
    }

    private void mergeBaseFields(RagHit target, RagHit incoming) {
        if ((target.getTitle() == null || target.getTitle().isEmpty()) && incoming.getTitle() != null) {
            target.setTitle(incoming.getTitle());
        }
        if ((target.getContent() == null || target.getContent().isEmpty()) && incoming.getContent() != null) {
            target.setContent(incoming.getContent());
        }
        if (incoming.getMetadata() != null && !incoming.getMetadata().isEmpty()) {
            target.getMetadata().putAll(incoming.getMetadata());
        }
    }

    private void assignRank(List<RagHit> results) {
        for (int i = 0; i < results.size(); i++) {
            RagHit hit = results.get(i);
            hit.setRank(i + 1);
            hit.getMetadata().put(RagRetrievalMetadataKeys.FINAL_RANK, i + 1);
            hit.getMetadata().put(RagRetrievalMetadataKeys.HIT_SOURCE, hit.getHitSource() == null ? null : hit.getHitSource().name());
            if (!hit.getMetadata().containsKey(RagRetrievalMetadataKeys.RERANK_SCORE) && hit.getScore() != null) {
                hit.getMetadata().put(RagRetrievalMetadataKeys.FUSION_SCORE, hit.getScore());
            }
            if (hit.getVectorScore() != null) {
                hit.getMetadata().put(RagRetrievalMetadataKeys.VECTOR_SCORE, hit.getVectorScore());
            }
            if (hit.getKeywordScore() != null) {
                hit.getMetadata().put(RagRetrievalMetadataKeys.KEYWORD_SCORE, hit.getKeywordScore());
            }
        }
    }

    private double nullSafeScore(Double value) {
        return value == null ? 0D : value;
    }
}
