package com.easyagents.rag.retrieval;

import com.easyagents.core.document.Document;

import java.util.HashMap;
import java.util.Map;

public class RagHit {

    private Object documentId;
    private String title;
    private String content;
    private Double score;
    private HitSource hitSource;
    private Double vectorScore;
    private Double keywordScore;
    private Integer rank;
    private Map<String, Object> metadata = new HashMap<String, Object>();

    public static RagHit fromDocument(Document document) {
        if (document == null) {
            return null;
        }
        RagHit hit = new RagHit();
        hit.setDocumentId(document.getId());
        hit.setTitle(document.getTitle());
        hit.setContent(document.getContent());
        hit.setScore(document.getScore());
        hit.setMetadata(document.getMetadataMap());

        Object hitSource = hit.getMetadata().get(RagRetrievalMetadataKeys.HIT_SOURCE);
        if (hitSource instanceof String) {
            hit.setHitSource(HitSource.valueOf(String.valueOf(hitSource)));
        }
        hit.setVectorScore(asDouble(hit.getMetadata().get(RagRetrievalMetadataKeys.VECTOR_SCORE)));
        hit.setKeywordScore(asDouble(hit.getMetadata().get(RagRetrievalMetadataKeys.KEYWORD_SCORE)));
        hit.setRank(asInteger(hit.getMetadata().get(RagRetrievalMetadataKeys.FINAL_RANK)));
        return hit;
    }

    public static RagHit fromDocument(Document document, HitSource defaultSource) {
        RagHit hit = fromDocument(document);
        if (hit != null && hit.getHitSource() == null) {
            hit.setHitSource(defaultSource);
        }
        if (hit != null && defaultSource == HitSource.VECTOR && hit.getVectorScore() == null) {
            hit.setVectorScore(document.getScore());
        }
        if (hit != null && defaultSource == HitSource.KEYWORD && hit.getKeywordScore() == null) {
            hit.setKeywordScore(document.getScore());
        }
        return hit;
    }

    public Document toDocument() {
        Document document = new Document();
        document.setId(documentId);
        document.setTitle(title);
        document.setContent(content);
        document.setScore(score);
        Map<String, Object> metadataMap = metadata == null
            ? new HashMap<String, Object>()
            : new HashMap<String, Object>(metadata);
        if (hitSource != null) {
            metadataMap.put(RagRetrievalMetadataKeys.HIT_SOURCE, hitSource.name());
        }
        if (vectorScore != null) {
            metadataMap.put(RagRetrievalMetadataKeys.VECTOR_SCORE, vectorScore);
        }
        if (keywordScore != null) {
            metadataMap.put(RagRetrievalMetadataKeys.KEYWORD_SCORE, keywordScore);
        }
        if (!metadataMap.containsKey(RagRetrievalMetadataKeys.FUSION_SCORE)
            && hitSource == HitSource.BOTH
            && score != null) {
            metadataMap.put(RagRetrievalMetadataKeys.FUSION_SCORE, score);
        }
        if (rank != null) {
            metadataMap.put(RagRetrievalMetadataKeys.FINAL_RANK, rank);
        }
        document.setMetadataMap(metadataMap);
        return document;
    }

    public RagHit copy() {
        RagHit copy = new RagHit();
        copy.setDocumentId(documentId);
        copy.setTitle(title);
        copy.setContent(content);
        copy.setScore(score);
        copy.setHitSource(hitSource);
        copy.setVectorScore(vectorScore);
        copy.setKeywordScore(keywordScore);
        copy.setRank(rank);
        copy.setMetadata(metadata);
        return copy;
    }

    public Object getDocumentId() {
        return documentId;
    }

    public void setDocumentId(Object documentId) {
        this.documentId = documentId;
    }

    public String getTitle() {
        return title;
    }

    public void setTitle(String title) {
        this.title = title;
    }

    public String getContent() {
        return content;
    }

    public void setContent(String content) {
        this.content = content;
    }

    public Double getScore() {
        return score;
    }

    public void setScore(Double score) {
        this.score = score;
    }

    public HitSource getHitSource() {
        return hitSource;
    }

    public void setHitSource(HitSource hitSource) {
        this.hitSource = hitSource;
    }

    public Double getVectorScore() {
        return vectorScore;
    }

    public void setVectorScore(Double vectorScore) {
        this.vectorScore = vectorScore;
    }

    public Double getKeywordScore() {
        return keywordScore;
    }

    public void setKeywordScore(Double keywordScore) {
        this.keywordScore = keywordScore;
    }

    public Integer getRank() {
        return rank;
    }

    public void setRank(Integer rank) {
        this.rank = rank;
    }

    public Map<String, Object> getMetadata() {
        return metadata;
    }

    public void setMetadata(Map<String, Object> metadata) {
        this.metadata = metadata == null ? new HashMap<String, Object>() : new HashMap<String, Object>(metadata);
    }

    private static Double asDouble(Object value) {
        if (value instanceof Number) {
            return ((Number) value).doubleValue();
        }
        if (value instanceof String) {
            try {
                return Double.valueOf((String) value);
            } catch (NumberFormatException ignore) {
                return null;
            }
        }
        return null;
    }

    private static Integer asInteger(Object value) {
        if (value instanceof Number) {
            return ((Number) value).intValue();
        }
        if (value instanceof String) {
            try {
                return Integer.valueOf((String) value);
            } catch (NumberFormatException ignore) {
                return null;
            }
        }
        return null;
    }
}
