# memory_manager.py - Cleaned version

import os
import shutil
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.docstore.document import Document

# --- Configuration ---
MEMORY_DIR = "memory"
INDEX_NAME = "faiss"
MODEL_NAME = "all-MiniLM-L6-v2"

class MemoryManager:
    def __init__(self):
        self.embeddings = HuggingFaceEmbeddings(model_name=MODEL_NAME)
        self.vector_store = self._load_or_create_vector_store()

    def reset_memory(self):
        """Removes the memory directory and re-initializes a new, empty index."""
        if os.path.exists(MEMORY_DIR):
            shutil.rmtree(MEMORY_DIR)
            os.makedirs(MEMORY_DIR, exist_ok=True)
            print("🧠 Memory reset successfully.")
            self.vector_store = self._create_new_index()

    def _load_or_create_vector_store(self):
        """Loads FAISS index or creates a new one, handling potential corruption."""
        index_path = os.path.join(MEMORY_DIR, f"{INDEX_NAME}.faiss")
        if os.path.exists(index_path):
            try:
                print("🧠 Loading existing memory from disk...")
                return FAISS.load_local(
                    folder_path=MEMORY_DIR, 
                    embeddings=self.embeddings,
                    index_name=INDEX_NAME,
                    allow_dangerous_deserialization=True
                )
            except Exception as e:
                print(f"⚠️ Error loading memory index: {e}. Rebuilding index.")
                shutil.rmtree(MEMORY_DIR)
                os.makedirs(MEMORY_DIR, exist_ok=True)
                return self._create_new_index()
        else:
            print("🧠 No existing memory found. Creating a new one.")
            return self._create_new_index()

    def _create_new_index(self):
        """Creates a fresh, empty FAISS index."""
        dummy_doc = [Document(page_content="Initial memory entry.")]
        # Note: If memory needs to be truly empty, use a small, persistent dummy doc
        # or handle an empty index creation if FAISS allows it. Keeping dummy for robustness.
        vs = FAISS.from_documents(dummy_doc, self.embeddings)
        vs.save_local(folder_path=MEMORY_DIR, index_name=INDEX_NAME)
        return vs

    def add_to_memory(self, text_to_add: str, metadata: dict):
        print(f"📝 Adding new memory: {text_to_add[:100]}...")
        doc = Document(page_content=text_to_add, metadata=metadata)
        self.vector_store.add_documents([doc])
        self.vector_store.save_local(folder_path=MEMORY_DIR, index_name=INDEX_NAME)

    def retrieve_relevant_memories(self, query: str, k: int = 5) -> list[Document]:
        print(f"🔍 Searching memory for: {query[:50]}...")
        return self.vector_store.similarity_search(query, k=k)

# --- Instantiate the class globally to satisfy 'from memory_manager import memory_manager' ---
memory_manager = MemoryManager()