package com.easyagents.rag.retrieval;

import org.junit.Assert;
import org.junit.Test;

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

public class RrfFusionStrategyTest {

    @Test
    public void shouldMarkDuplicateHitsAsBothAndRankFirst() {
        RrfFusionStrategy strategy = new RrfFusionStrategy();

        RagHit vectorOnly = hit("1", 0.91D, HitSource.VECTOR);
        RagHit bothVector = hit("2", 0.88D, HitSource.VECTOR);
        RagHit bothKeyword = hit("2", 5.2D, HitSource.KEYWORD);
        RagHit keywordOnly = hit("3", 4.8D, HitSource.KEYWORD);

        List<RagHit> fused = strategy.fuse(
            Arrays.asList(vectorOnly, bothVector),
            Arrays.asList(bothKeyword, keywordOnly),
            10
        );

        Assert.assertEquals(3, fused.size());
        Assert.assertEquals("2", String.valueOf(fused.get(0).getDocumentId()));
        Assert.assertEquals(HitSource.BOTH, fused.get(0).getHitSource());
        Assert.assertEquals("1", String.valueOf(fused.get(1).getDocumentId()));
        Assert.assertEquals("3", String.valueOf(fused.get(2).getDocumentId()));
    }

    private RagHit hit(String id, double score, HitSource source) {
        RagHit hit = new RagHit();
        hit.setDocumentId(id);
        hit.setContent("doc-" + id);
        hit.setScore(score);
        hit.setHitSource(source);
        if (source == HitSource.VECTOR) {
            hit.setVectorScore(score);
        } else if (source == HitSource.KEYWORD) {
            hit.setKeywordScore(score);
        }
        return hit;
    }
}
