"""
============================================================
PDF Ingestion - يبني ChromaDB
============================================================
Usage:
  python ingest.py                  # uses default paths
  python ingest.py --data ./pdfs    # custom data folder
"""
import os
import argparse
import shutil
from langchain_community.document_loaders import PyPDFDirectoryLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from dotenv import load_dotenv

from config import CHROMA_DB_PATH, PDF_DOCS_DIR, RAG_CONFIG

load_dotenv()


def build_vector_database(data_dir: str = None, rebuild: bool = False):
    """بناء/تحديث الـ ChromaDB من PDFs"""
    if data_dir is None:
        data_dir = PDF_DOCS_DIR

    if not os.path.exists(data_dir):
        print(f"✗ Folder '{data_dir}' does not exist.")
        print("  Create it and add your PDF files, then run again.")
        return False

    pdfs = [f for f in os.listdir(data_dir) if f.lower().endswith('.pdf')]
    if not pdfs:
        print(f"✗ No PDF files in '{data_dir}'.")
        return False

    print(f"📂 Loading {len(pdfs)} PDF(s) from '{data_dir}'...")
    loader = PyPDFDirectoryLoader(data_dir)
    documents = loader.load()

    if not documents:
        print("✗ Loaded 0 pages.")
        return False

    print(f"✓ Loaded {len(documents)} pages. Splitting into chunks...")

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=RAG_CONFIG['chunk_size'],
        chunk_overlap=RAG_CONFIG['chunk_overlap']
    )
    chunks = splitter.split_documents(documents)
    print(f"  → {len(chunks)} chunks created.")

    print(f"⚙ Generating embeddings ({RAG_CONFIG['embedding_model']})...")
    embeddings = HuggingFaceEmbeddings(model_name=RAG_CONFIG['embedding_model'])

    if rebuild and os.path.exists(CHROMA_DB_PATH):
        print(f"🗑 Removing existing DB at '{CHROMA_DB_PATH}' (rebuild flag set)")
        shutil.rmtree(CHROMA_DB_PATH)

    if os.path.exists(CHROMA_DB_PATH) and not rebuild:
        print(f"⚠ Existing DB at '{CHROMA_DB_PATH}' — chunks will be ADDED on top.")
        print("  Use --rebuild for a clean rebuild.")

    Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=CHROMA_DB_PATH
    )

    print(f"\n✓ Vector DB ready: {len(documents)} pages → {len(chunks)} chunks at '{CHROMA_DB_PATH}'")
    return True


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', default=None, help='PDF directory')
    parser.add_argument('--rebuild', action='store_true', help='Wipe existing DB first')
    args = parser.parse_args()
    build_vector_database(args.data, args.rebuild)