# rag/db/initializer.py
import faiss
import numpy as np
from huggingface_hub import hf_hub_download
from config import HF_DS_REPO_ID, HF_INDEX_FILE, HF_IDS_FILE
from modules.retriever import set_index
from modules import corpus

_vector_ids = None

def _load_index_in_memory():
    """HF Hub에서 인덱스/ID 매핑을 받아 메모리에 로드"""
    index_path = hf_hub_download(repo_id=HF_DS_REPO_ID, filename=HF_INDEX_FILE, repo_type="dataset")
    ids_path = hf_hub_download(repo_id=HF_DS_REPO_ID, filename=HF_IDS_FILE, repo_type="dataset")
    index = faiss.read_index(index_path)
    set_index(index)
    global _vector_ids
    _vector_ids = np.load(ids_path, allow_pickle=True)

def get_vector_ids():
    global _vector_ids
    return _vector_ids

def initialize_dbs():
    # 1) 코퍼스 준비 (최초 1회만 다운로드)
    corpus.prepare_corpus()
    # 2) 인덱스/ID 매핑 메모리 로드
    _load_index_in_memory()

def force_update():
    _load_index_in_memory()