import numpy as np
from sklearn.metrics.pairwise import cosine_similarity


class SemanticCache:
    def __init__(self, similarity_threshold=0.70):
        self.cache = []
        self.similarity_threshold = similarity_threshold
        self.hit_count = 0
        self.miss_count = 0

    def search(self, query_embedding):
        """
        Check if a similar query exists in cache
        """
        if len(self.cache) == 0:
            self.miss_count += 1
            return None

        embeddings = [item["embedding"] for item in self.cache]

        similarities = cosine_similarity(query_embedding, embeddings)[0]

        best_idx = np.argmax(similarities)
        best_score = similarities[best_idx]

        if best_score >= self.similarity_threshold:
            self.hit_count += 1
            return {
                "result": self.cache[best_idx]["result"],
                "matched_query": self.cache[best_idx]["query"],
                "similarity_score": float(best_score)
            }

        self.miss_count += 1
        return None

    def add(self, query, embedding, result):
        """
        Store query and result in cache
        """
        self.cache.append({
            "query": query,
            "embedding": embedding[0],
            "result": result
        })

    def stats(self):
        total = self.hit_count + self.miss_count
        hit_rate = self.hit_count / total if total > 0 else 0

        return {
            "total_entries": len(self.cache),
            "hit_count": self.hit_count,
            "miss_count": self.miss_count,
            "hit_rate": hit_rate
        }

    def clear(self):
        self.cache = []
        self.hit_count = 0
        self.miss_count = 0