"""Lab 2 - Bangun indeks ChromaDB dari semua PDF di folder dokumen/.

Pakai:
    python indeks.py                 # ukuran 800, overlap 100 (koleksi: pedoman_800)
    python indeks.py 500 1000        # beberapa ukuran sekaligus (overlap 15%)
Tiap ukuran menjadi koleksi sendiri (pedoman_500, pedoman_1000, ...) untuk Lab 3.
Aman dijalankan ulang (upsert).
"""
import sys
from umum import (FOLDER_DOKUMEN, baca_pdf_per_halaman, potong, nama_koleksi,
                  ambil_embedding, buka_klien, buat_koleksi)

if len(sys.argv) > 1:
    daftar = [(int(a), int(int(a) * 0.15)) for a in sys.argv[1:]]
else:
    daftar = [(800, 100)]

pdfs = sorted(FOLDER_DOKUMEN.glob("*.pdf"))
if not pdfs:
    sys.exit("Tidak ada PDF di folder dokumen/. Letakkan 2-3 PDF berbasis teks di sana.")

print("Memuat model embedding (pertama kali akan mengunduh ratusan MB)...")
emb = ambil_embedding()
klien = buka_klien()

for ukuran, tumpang in daftar:
    koleksi = buat_koleksi(klien, nama_koleksi(ukuran), emb)
    print(f"\n== Koleksi {koleksi.name} (ukuran={ukuran}, overlap={tumpang}) ==")
    for pdf in pdfs:
        halaman = baca_pdf_per_halaman(pdf)
        if not halaman:
            print(f"{pdf.name}: TEKS KOSONG (PDF scan?) - dilewati")
            continue
        ids, docs, metas = [], [], []
        n = 0
        for no_hal, teks in halaman:                 # potong per halaman agar metadata halaman akurat
            for c in potong(teks, ukuran, tumpang):
                ids.append(f"{pdf.stem}-{ukuran}-{n}")
                docs.append(c)
                metas.append({"sumber": pdf.name, "halaman": no_hal, "chunk": n})
                n += 1
        for i in range(0, len(ids), 100):             # kirim bertahap
            koleksi.upsert(ids=ids[i:i+100], documents=docs[i:i+100], metadatas=metas[i:i+100])
        print(f"{pdf.name} -> {len(ids)} chunk")
    print(f"Total di koleksi: {koleksi.count()} chunk")