import os
import chromadb
from chromadb.utils import embedding_functions
from pypdf import PdfReader
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 1. Persiapan Dokumen
pdf_path = "dokumen_akademik.pdf" # Ganti dengan nama file PDF yang disiapkan
print(f"Membaca file: {pdf_path}")

reader = PdfReader('Buku Pedoman Akademik.pdf')
raw_text = ""
for page in reader.pages:
    # Menggabungkan seluruh teks dari PDF
    raw_text += page.extract_text() + "\n"

# 2. Setup Embedding Function
# Menggunakan model multilingual lokal agar performa pencarian untuk Bahasa Indonesia lebih akurat
print("Menyiapkan model embedding (ini akan mengunduh model pada run pertama)...")
emb_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="paraphrase-multilingual-MiniLM-L12-v2"
)

# 3. Inisialisasi ChromaDB Client
# PersistentClient akan membuat folder 'chroma_db_data' untuk menyimpan vektor secara lokal
client = chromadb.PersistentClient(path="./chroma_db_data")

# Eksperimen 3 ukuran chunk sesuai instruksi modul
chunk_sizes = [500, 1000, 1500]

for size in chunk_sizes:
    print(f"\n--- Memproses ukuran chunk: {size} karakter ---")
    collection_name = f"rag_koleksi_chunk_{size}"
    
    # Hapus collection jika sebelumnya sudah ada saat di-run ulang
    try:
        client.delete_collection(name=collection_name)
    except Exception:
        pass
        
    # Buat collection baru di vector DB
    collection = client.create_collection(
        name=collection_name, 
        embedding_function=emb_fn
    )
    
    # 4. Setup Text Splitter
    # Overlap disetel 10% dari ukuran chunk untuk mencegah konteks terpotong mendadak
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=size,
        chunk_overlap=int(size * 0.1), 
        length_function=len,
    )
    
    # Potong teks menjadi list of chunks
    chunks = text_splitter.split_text(raw_text)
    print(f"Total chunk yang dihasilkan: {len(chunks)} chunks")
    
    # 5. Persiapkan ID dan Metadata
    ids = [f"id_chunk_{size}_{i}" for i in range(len(chunks))]
    metadatas = [{"source": pdf_path, "chunk_size": size, "chunk_index": i} for i in range(len(chunks))]
    
    # 6. Masukkan Vektor, Teks, dan Metadata ke Vector Database
    collection.add(
        documents=chunks,
        metadatas=metadatas,
        ids=ids
    )
    print(f"Data chunk {size} berhasil di-index ke ChromaDB!")

# ==========================================
# CONTOH PENCARIAN (QUERY)
# ==========================================
print("\n--- Menguji Pencarian (Query Retrieval) ---")
query_text = "Apa syarat untuk mengajukan cuti akademik?" # Sesuaikan dengan isi PDF yang di-load

# Kita uji pencarian di collection yang ukuran chunk-nya 1000
test_collection = client.get_collection(name="rag_koleksi_chunk_1000", embedding_function=emb_fn)

# Retrieve Top-K (K=2)
results = test_collection.query(
    query_texts=[query_text],
    n_results=2
)

print(f"Pertanyaan: '{query_text}'\n")

# Hasil dibungkus di dalam list of lists oleh chromadb, maka kita akses index [0]
for i, (doc, dist) in enumerate(zip(results['documents'][0], results['distances'][0])):
    print(f"--- Top {i+1} (Jarak Vektor: {dist:.4f}) ---")
    print(f"{doc}\n")