import os
import faiss
import threading

class UserIndexManager:
    def __init__(self, base_dir="index_data/"):
        self.base_dir = base_dir
        self.index_map = {}  # user_id -> faiss.Index
        self.index_lock = threading.Lock()
        self._load_all_indexes()

    def _load_all_indexes(self):
        for fname in os.listdir(self.base_dir):
            if fname.endswith(".index"):
                user_id = fname.replace(".index", "")
                self._load_single_index(user_id)

    def _load_single_index(self, user_id):
        path = os.path.join(self.base_dir, f"{user_id}.index")
        if os.path.exists(path):
            index = faiss.read_index(path)
            self.index_map[user_id] = index
            print(f"[INIT] Loaded index for user: {user_id}")

    def get_index(self, user_id):
        with self.index_lock:
            if user_id not in self.index_map:
                raise FileNotFoundError(f"No FAISS index loaded for user: {user_id}")
            return self.index_map[user_id]

    def update_index(self, user_id):
        """重新加载用户索引文件"""
        path = os.path.join(self.base_dir, f"{user_id}.index")
        if not os.path.exists(path):
            raise FileNotFoundError(f"No index file found for user: {user_id}")
        with self.index_lock:
            self.index_map[user_id] = faiss.read_index(path)
            print(f"[UPDATE] Index for user {user_id} has been hot-reloaded.")