package com.easyagents.rag.retrieval;

import com.easyagents.core.document.Document;
import com.easyagents.core.model.rerank.RerankModel;

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

public class RagRetrievalExecutor {

    private final VectorRetriever vectorRetriever;
    private final KeywordRetriever keywordRetriever;
    private final FusionStrategy fusionStrategy;

    public RagRetrievalExecutor(VectorRetriever vectorRetriever,
                                KeywordRetriever keywordRetriever,
                                FusionStrategy fusionStrategy) {
        this.vectorRetriever = vectorRetriever;
        this.keywordRetriever = keywordRetriever;
        this.fusionStrategy = fusionStrategy == null ? new RrfFusionStrategy() : fusionStrategy;
    }

    public RagRetrievalResult retrieve(RagQuery query) {
        RagQuery effectiveQuery = query == null ? new RagQuery() : query;
        RetrievalMode mode = effectiveQuery.getRetrievalMode() == null
            ? RetrievalMode.HYBRID
            : effectiveQuery.getRetrievalMode();
        if (mode == RetrievalMode.VECTOR) {
            return buildResult(normalizeHits(vectorRetriever == null
                ? Collections.<RagHit>emptyList()
                : vectorRetriever.retrieve(effectiveQuery), HitSource.VECTOR), effectiveQuery.getTopK());
        }
        if (mode == RetrievalMode.KEYWORD) {
            return buildResult(normalizeHits(keywordRetriever == null
                ? Collections.<RagHit>emptyList()
                : keywordRetriever.retrieve(effectiveQuery), HitSource.KEYWORD), effectiveQuery.getTopK());
        }

        CompletableFuture<List<RagHit>> vectorFuture = CompletableFuture.supplyAsync(new java.util.function.Supplier<List<RagHit>>() {
            @Override
            public List<RagHit> get() {
                return vectorRetriever == null ? Collections.<RagHit>emptyList() : vectorRetriever.retrieve(effectiveQuery);
            }
        });
        CompletableFuture<List<RagHit>> keywordFuture = CompletableFuture.supplyAsync(new java.util.function.Supplier<List<RagHit>>() {
            @Override
            public List<RagHit> get() {
                return keywordRetriever == null ? Collections.<RagHit>emptyList() : keywordRetriever.retrieve(effectiveQuery);
            }
        });

        List<RagHit> vectorHits = normalizeHits(vectorFuture.join(), HitSource.VECTOR);
        List<RagHit> keywordHits = normalizeHits(keywordFuture.join(), HitSource.KEYWORD);
        return buildResult(fusionStrategy.fuse(vectorHits, keywordHits, effectiveQuery.getTopK()), effectiveQuery.getTopK());
    }

    public RagRetrievalResult rerank(String query, List<RagHit> hits, RerankModel rerankModel, int topK) {
        if (rerankModel == null || hits == null || hits.isEmpty()) {
            return buildResult(hits, topK);
        }
        List<Document> rerankInput = new ArrayList<Document>(hits.size());
        Map<String, RagHit> hitMap = new LinkedHashMap<String, RagHit>();
        for (RagHit hit : hits) {
            if (hit == null || hit.getDocumentId() == null) {
                continue;
            }
            rerankInput.add(hit.toDocument());
            hitMap.put(String.valueOf(hit.getDocumentId()), hit.copy());
        }
        List<Document> rerankedDocuments = rerankModel.rerank(query, rerankInput);
        List<RagHit> rerankedHits = new ArrayList<RagHit>();
        if (rerankedDocuments != null) {
            for (Document rerankedDocument : rerankedDocuments) {
                if (rerankedDocument == null || rerankedDocument.getId() == null) {
                    continue;
                }
                RagHit original = hitMap.get(String.valueOf(rerankedDocument.getId()));
                if (original == null) {
                    continue;
                }
                original.setContent(rerankedDocument.getContent());
                original.setTitle(rerankedDocument.getTitle());
                original.setScore(rerankedDocument.getScore());
                if (rerankedDocument.getScore() != null) {
                    original.getMetadata().put(RagRetrievalMetadataKeys.RERANK_SCORE, rerankedDocument.getScore());
                }
                rerankedHits.add(original);
            }
        }
        sortByScore(rerankedHits);
        assignRank(rerankedHits);
        return buildResult(rerankedHits, topK);
    }

    private RagRetrievalResult buildResult(List<RagHit> hits, int topK) {
        List<RagHit> effectiveHits = hits == null ? new ArrayList<RagHit>() : new ArrayList<RagHit>(hits);
        if (topK > 0 && effectiveHits.size() > topK) {
            effectiveHits = new ArrayList<RagHit>(effectiveHits.subList(0, topK));
        }
        assignRank(effectiveHits);
        RagRetrievalResult result = new RagRetrievalResult();
        result.setHits(effectiveHits);
        result.setTotal(effectiveHits.size());
        return result;
    }

    private List<RagHit> normalizeHits(List<RagHit> hits, HitSource defaultSource) {
        List<RagHit> normalized = new ArrayList<RagHit>();
        if (hits == null) {
            return normalized;
        }
        for (RagHit hit : hits) {
            if (hit == null || hit.getDocumentId() == null) {
                continue;
            }
            RagHit copy = hit.copy();
            if (copy.getHitSource() == null) {
                copy.setHitSource(defaultSource);
            }
            if (defaultSource == HitSource.VECTOR && copy.getVectorScore() == null) {
                copy.setVectorScore(copy.getScore());
            }
            if (defaultSource == HitSource.KEYWORD && copy.getKeywordScore() == null) {
                copy.setKeywordScore(copy.getScore());
            }
            normalized.add(copy);
        }
        assignRank(normalized);
        return normalized;
    }

    private void sortByScore(List<RagHit> hits) {
        Collections.sort(hits, new Comparator<RagHit>() {
            @Override
            public int compare(RagHit left, RagHit right) {
                return Double.compare(right.getScore() == null ? 0D : right.getScore(),
                    left.getScore() == null ? 0D : left.getScore());
            }
        });
    }

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