# create_vector_store.py — Rebuild vector store from markdown data files
import os
import shutil
from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from dotenv import load_dotenv

# --- Project Settings ---
DATA_PATH = os.path.join(os.path.dirname(__file__), "university_data_arabic")
DB_PATH = os.path.join(os.path.dirname(__file__), "vectorstores", "db")
EMBEDDING_MODEL_NAME = "intfloat/multilingual-e5-small"
HF_CACHE_DIR = os.path.join(os.path.dirname(__file__), ".cache", "huggingface")


def build_vector_store():
    """
    Builds the ChromaDB vector store from markdown files.
    
    Workflow:
    1. Loads all .md files from university_data_arabic/
    2. Splits them into ~1000 char chunks with 200 char overlap
    3. Embeds using intfloat/multilingual-e5-small (local model, cached on disk)
    4. Persists the vector store to vectorstores/db/
    
    The embedding model (~470MB) is downloaded only once and cached in .cache/
    """
    load_dotenv()

    if not os.path.exists(DATA_PATH):
        print(f"❌ Data folder '{DATA_PATH}' not found.")
        print(f"   Create it and add your .md files there.")
        return

    # Clean old vector store to ensure fresh data
    if os.path.exists(DB_PATH):
        print(f"🗑️  Removing old vector store at {DB_PATH}")
        shutil.rmtree(DB_PATH)

    # 1. Load all .md files
    print(f"📂 Loading files from: {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("❌ No .md files found. Add markdown files to the data directory.")
        return

    print(f"   Found {len(documents)} documents")

    # 2. Split into chunks
    print("✂️  Splitting documents into chunks...")
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200
    )
    chunked_documents = text_splitter.split_documents(documents)
    print(f"   Created {len(chunked_documents)} chunks")

    # 3. Initialize embedding model
    print(f"🧠 Loading embedding model: {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. Build and persist ChromaDB
    print(f"💾 Building vector store at: {DB_PATH}...")
    vector_store = Chroma.from_documents(
        documents=chunked_documents,
        embedding=embedding_model,
        persist_directory=DB_PATH
    )

    print("-" * 45)
    print("✅ Vector store built successfully!")
    print(f"   Documents: {len(documents)}")
    print(f"   Chunks: {len(chunked_documents)}")
    print(f"   Location: {DB_PATH}")
    print("-" * 45)


if __name__ == "__main__":
    build_vector_store()