import os
import faiss
import pickle
from sentence_transformers import SentenceTransformer
from logger import logger

class MemoryVectorStore:
    def __init__(self, path="memory_index"):
        self.path = path
        self.model = SentenceTransformer("all-MiniLM-L6-v2")
        self.dim = 384
        self.index = faiss.IndexFlatL2(self.dim)
        self.memory = []
        self.load()

    def load(self):
        if os.path.exists(self.path + ".pkl"):
            with open(self.path + ".pkl", "rb") as f:
                data = pickle.load(f)
                self.memory = data["memory"]
                self.index = data["index"]
            logger.info("✅ Memory index loaded")

    def save(self):
        with open(self.path + ".pkl", "wb") as f:
            pickle.dump({"memory": self.memory, "index": self.index}, f)

    def add(self, text):
        vec = self.model.encode([text])
        self.index.add(vec)
        self.memory.append(text)
        self.save()

    def search(self, query, k=3):
        vec = self.model.encode([query])
        D, I = self.index.search(vec, k)
        return [self.memory[i] for i in I[0] if i < len(self.memory)]