package com.ankurm.semanticcache;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;

import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.redis.RedisVectorStore;
import redis.clients.jedis.RedisClient;

/**
 * A question-to-answer cache over a Redis vector index. The key is the question's embedding, the value is the
 * answer, and a lookup is a nearest-neighbour search with a similarity threshold.
 */
public final class SemanticCache {

    /** A stored answer and how close the stored question was to the one asked. */
    public record Hit(String answer, String storedQuestion, double score) { }

    public static final String INDEX = "semantic_cache";
    public static final String PREFIX = "sc:";

    private final RedisVectorStore store;
    private final RedisClient jedis;
    private volatile double threshold;

    public SemanticCache(RedisClient jedis, EmbeddingModel embeddings, double threshold) {
        this.jedis = jedis;
        this.threshold = threshold;
        try {
            jedis.ftDropIndex(INDEX);
        } catch (Exception ignored) {
            // no index yet
        }
        this.store = RedisVectorStore.builder(jedis, embeddings)
                .indexName(INDEX)
                .prefix(PREFIX)
                .vectorAlgorithm(RedisVectorStore.Algorithm.HNSW)
                .distanceMetric(RedisVectorStore.DistanceMetric.COSINE)
                // Only declared metadata fields are filterable, and only declared ones come back on a search result:
                // without text("answer") a hit has the stored question but no answer.
                .metadataFields(RedisVectorStore.MetadataField.tag("tenant"), RedisVectorStore.MetadataField.text("answer"))
                .initializeSchema(true)
                .build();
        this.store.afterPropertiesSet();
    }

    public void threshold(double t) { this.threshold = t; }

    public double threshold() { return threshold; }

    public Optional<Hit> lookup(String tenant, String question) {
        return lookup(tenant, question, threshold);
    }

    public Optional<Hit> lookup(String tenant, String question, double minScore) {
        List<Document> found = store.similaritySearch(SearchRequest.builder()
                .query(question)
                .topK(1)
                .similarityThreshold(minScore)
                .filterExpression("tenant == '" + tenant + "'")
                .build());
        return found.stream().findFirst().map(d -> new Hit(
                (String) d.getMetadata().get("answer"), d.getText(), d.getScore()));
    }

    /** The nearest stored question whatever its score, for measuring rather than serving. */
    public Optional<Hit> nearest(String tenant, String question) {
        return lookup(tenant, question, 0.0);
    }

    public String put(String tenant, String question, String answer) {
        String id = UUID.randomUUID().toString();
        store.add(List.of(new Document(id, question, Map.of("tenant", tenant, "answer", answer))));
        return id;
    }

    /** Redis expires the whole document; the index drops it with the key. */
    public void expireAfterSeconds(String id, long seconds) {
        jedis.expire(PREFIX + id, seconds);
    }

    public long size() {
        return jedis.ftInfo(INDEX).get("num_docs") instanceof Number n ? n.longValue()
                : Long.parseLong(String.valueOf(jedis.ftInfo(INDEX).get("num_docs")));
    }
}