# create_vector_store.py
from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_community.embeddings import HuggingFaceEmbeddings
from dotenv import load_dotenv
import os
import shutil

# --- إعدادات المشروع ---
DATA_PATH = "university_data_arabic"
DB_PATH = "vectorstores/db/"
# نستخدم تضمينات جوجل المستضافة لتقليل حجم الحزم وعدم تحميل نماذج محلياً
# EMBEDDING_MODEL_NAME = "models/text-embedding-004"
EMBEDDING_MODEL_NAME = "intfloat/multilingual-e5-small"
HF_CACHE_DIR = "/app/.cache/huggingface" # could be deleted if not using local models
# # ---------------------

def build_vector_store():
    """
    يقوم هذا السكربت ببناء قاعدة البيانات المعرفية (Vector Store)
    من ملفات الماركداون الموجودة في مجلد البيانات.
    """
    # تحميل المتغيرات البيئية (مثل GOOGLE_API_KEY)
    load_dotenv()
    # التأكد من وجود مجلد البيانات
    if not os.path.exists(DATA_PATH):
        print(f"خطأ: مجلد البيانات '{DATA_PATH}' غير موجود. يرجى إنشاؤه ووضع ملفات البيانات بداخله.")
        return

    # تنظيف قاعدة البيانات القديمة إذا كانت موجودة لضمان التحديث
    if os.path.exists(DB_PATH):
        print(f"يتم حذف قاعدة البيانات القديمة في {DB_PATH}")
        shutil.rmtree(DB_PATH)

    # 1. تحميل كل ملفات .md من مجلد البيانات
    print(f"يتم تحميل الملفات من المسار: {DATA_PATH}...")
    loader = DirectoryLoader(DATA_PATH, glob="**/*.md", loader_cls=TextLoader, loader_kwargs={'encoding': 'utf-8'}, show_progress=True)
    documents = loader.load()
    if not documents:
        print("لم يتم العثور على أي ملفات. يرجى التأكد من وجود ملفات .md في مجلد البيانات.")
        return

    # 2. تقسيم المستندات إلى أجزاء (chunks) أصغر
    print("يتم تقسيم المستندات إلى أجزاء أصغر...")
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    chunked_documents = text_splitter.split_documents(documents)
    
    # 3. إنشاء دالة التضمين باستخدام واجهة جوجل (عن بعد)
    # لا يتم تحميل أي نموذج محلياً، مما يقلل المساحة المطلوبة بشكل كبير
    # print(f"يتم تهيئة تضمينات جوجل: {EMBEDDING_MODEL_NAME}...")
    # embedding_model = GoogleGenerativeAIEmbeddings(
    #     model=EMBEDDING_MODEL_NAME
    # )

    #print(f"يتم تهيئة تضمينات هاجنج فيس: {EMBEDDING_MODEL_NAME} ...")
    # embedding_model = HuggingFaceEmbeddings(
    #         model_name=EMBEDDING_MODEL_NAME,
    #         model_kwargs={'device': 'cpu'},
    #         encode_kwargs={'normalize_embeddings': True}
    #     )
    print(f"يتم تهيئة تضمينات هاجنج فيس: {EMBEDDING_MODEL_NAME} ...")
    embedding_model = HuggingFaceEmbeddings(
        model_name=EMBEDDING_MODEL_NAME,
        cache_folder=HF_CACHE_DIR,
        model_kwargs={'device': 'cpu'},
        encode_kwargs={'normalize_embeddings': True}
    )



    # 4. إنشاء قاعدة بيانات ChromaDB وحفظها على الجهاز
    print(f"يتم إنشاء قاعدة البيانات المعرفية في المسار: {DB_PATH}...")
    vector_store = Chroma.from_documents(
        documents=chunked_documents,
        embedding=embedding_model,
        persist_directory=DB_PATH
    )
    print("-----------------------------------------")
    print("✅ تم إنشاء قاعدة البيانات بنجاح! ✅")
    print("-----------------------------------------")


if __name__ == "__main__":
    build_vector_store()