import chromadb
import json
import argparse
from FlagEmbedding import FlagModel

# 初始化嵌入模型
embedding_model = FlagModel(
    'chuxin-llm/Chuxin-Embedding',
    query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：",
    use_fp16=True  # 使用半精度計算，加速推理
)

def get_chuxin_embedding(text):
    """生成文本的嵌入向量"""
    return embedding_model.encode([text])[0].tolist()

def load_data_to_chromadb(json_path, chromadb_path, collection_name):
    """將 JSON 數據載入到 ChromaDB"""
    # 讀取 JSON 文件
    with open(json_path, "r", encoding="utf-8") as f:
        documents = json.load(f)

    # 初始化 ChromaDB client
    chroma_client = chromadb.PersistentClient(path=chromadb_path)
    
    # 創建或獲取 Collection
    collection = chroma_client.get_or_create_collection(name=collection_name, embedding_function=None)
    
    # 解析 JSON 並添加到 ChromaDB
    existing_ids = set(collection.get()["ids"])  # 取得目前所有的 ID

    for doc in documents:
        embedding = get_chuxin_embedding(doc["text"])
        if doc["id"] in existing_ids:
            # 更新已存在的記錄
            collection.update(
                ids=[doc["id"]],
                documents=[doc["text"]],
                embeddings=[embedding],
                metadatas=[doc["metadata"]]
            )
            print(f"更新 {doc['id']} 的內容")
        else:
            # 新增新的記錄
            collection.add(
                ids=[doc["id"]],
                documents=[doc["text"]],
                embeddings=[embedding],
                metadatas=[doc["metadata"]]
            )
            print(f"新增 {doc['id']} 的內容")

    
    print(f"成功載入 {collection.count()} 筆資料到 ChromaDB")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="將 JSON 數據載入到 ChromaDB")
    parser.add_argument("json_path", type=str, help="JSON 文件的路徑")
    parser.add_argument("chromadb_path", type=str, help="ChromaDB 的存儲地址")
    parser.add_argument("collection_name", type=str, help="ChromaDB Collection 的名稱")
    
    args = parser.parse_args()
    
    load_data_to_chromadb(args.json_path, args.chromadb_path, args.collection_name)