package com.easyagents.rag.retrieval;

import com.easyagents.core.document.Document;
import com.easyagents.core.model.rerank.RerankModel;
import com.easyagents.core.model.rerank.RerankOptions;
import org.junit.Assert;
import org.junit.Test;

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

public class RagRetrievalExecutorTest {

    @Test
    public void shouldOnlyCallVectorRetrieverInVectorMode() {
        final List<String> invoked = new ArrayList<String>();
        RagRetrievalExecutor executor = new RagRetrievalExecutor(
            new VectorRetriever() {
                @Override
                public List<RagHit> retrieve(RagQuery query) {
                    invoked.add("vector");
                    return Arrays.asList(hit("1", "vector", 0.9D, HitSource.VECTOR));
                }
            },
            new KeywordRetriever() {
                @Override
                public List<RagHit> retrieve(RagQuery query) {
                    invoked.add("keyword");
                    return Arrays.asList(hit("2", "keyword", 5D, HitSource.KEYWORD));
                }
            },
            new RrfFusionStrategy()
        );

        RagQuery query = new RagQuery();
        query.setRetrievalMode(RetrievalMode.VECTOR);
        query.setTopK(10);

        RagRetrievalResult result = executor.retrieve(query);

        Assert.assertEquals(Arrays.asList("vector"), invoked);
        Assert.assertEquals(1, result.getHits().size());
        Assert.assertEquals(HitSource.VECTOR, result.getHits().get(0).getHitSource());
    }

    @Test
    public void shouldRerankAfterHybridFusionAndKeepHitSource() {
        RagRetrievalExecutor executor = new RagRetrievalExecutor(
            new VectorRetriever() {
                @Override
                public List<RagHit> retrieve(RagQuery query) {
                    return Arrays.asList(
                        hit("1", "alpha", 0.9D, HitSource.VECTOR),
                        hit("2", "beta", 0.8D, HitSource.VECTOR)
                    );
                }
            },
            new KeywordRetriever() {
                @Override
                public List<RagHit> retrieve(RagQuery query) {
                    return Arrays.asList(
                        hit("2", "beta", 4.1D, HitSource.KEYWORD),
                        hit("3", "gamma", 3.9D, HitSource.KEYWORD)
                    );
                }
            },
            new RrfFusionStrategy()
        );

        RagQuery query = new RagQuery();
        query.setRetrievalMode(RetrievalMode.HYBRID);
        query.setTopK(10);

        RagRetrievalResult retrieved = executor.retrieve(query);
        RagRetrievalResult reranked = executor.rerank("query", retrieved.getHits(), new ReverseRerankModel(), 10);

        Assert.assertEquals("3", String.valueOf(reranked.getHits().get(0).getDocumentId()));
        Assert.assertEquals("2", String.valueOf(reranked.getHits().get(2).getDocumentId()));
        Assert.assertEquals(HitSource.BOTH, reranked.getHits().get(2).getHitSource());
    }

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

    private static class ReverseRerankModel implements RerankModel {

        @Override
        public List<Document> rerank(String query, List<Document> documents, RerankOptions options) {
            List<Document> result = new ArrayList<Document>();
            double score = documents.size();
            for (int i = documents.size() - 1; i >= 0; i--) {
                Document document = documents.get(i);
                document.setScore(score--);
                result.add(document);
            }
            return result;
        }
    }
}
