import os
import json
import numpy as np

class MemoryStore:
    def __init__(self, path="/mnt/data/memory.json"):
        self.path = path
        os.makedirs(os.path.dirname(self.path), exist_ok=True)  # ensure folder exists
        self.memories = []
        self.load()

    def add_memory(self, text, vector):
        # avoid duplicates based on cosine similarity
        if self.memories:
            sims = [self.cosine_similarity(vector, np.array(m["vector"])) for m in self.memories]
            if max(sims) > 0.95:
                return "Duplicate detected — memory already stored."
        self.memories.append({"text": text, "vector": vector.tolist()})
        self.save()
        return "Saved new memory."

    def retrieve_relevant(self, query_vector, top_k=5):
        if not self.memories:
            return []
        scored = [
            (m["text"], self.cosine_similarity(query_vector, np.array(m["vector"])))
            for m in self.memories
        ]
        return sorted(scored, key=lambda x: x[1], reverse=True)[:top_k]

    def save(self):
        os.makedirs(os.path.dirname(self.path), exist_ok=True)
        with open(self.path, "w") as f:
            json.dump(self.memories, f)

    def load(self):
        if os.path.exists(self.path):
            with open(self.path, "r") as f:
                self.memories = json.load(f)
        else:
            self.memories = []

    @staticmethod
    def cosine_similarity(a, b):
        return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))