from pathlib import Path
from pypdf import PdfReader
import chromadb
from chromadb.utils import embedding_functions

def baca_pdf(path):
    r = PdfReader(path)
    return "\n".join((p.extract_text() or "") for p in r.pages)

def potong(teks, ukuran=500, tumpang=75):
    hasil, mulai = [], 0
    while mulai < len(teks):
        hasil.append(teks[mulai:mulai + ukuran])
        mulai += ukuran - tumpang
    return [h for h in hasil if h.strip()]

print("Menyiapkan model embedding (unduh jika belum ada)...")
# Model multilingual agar bisa bahasa Indonesia
emb = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="paraphrase-multilingual-MiniLM-L12-v2")

client = chromadb.PersistentClient(path="./chroma_db")
koleksi = client.get_or_create_collection(
    name="pedoman", embedding_function=emb,
    metadata={"hnsw:space": "cosine"})

# Proses tiap PDF
for pdf in Path("dokumen").glob("*.pdf"):
    teks = baca_pdf(pdf)
    chunks = potong(teks, ukuran=500, tumpang=75) # Pakai ukuran 500
    koleksi.upsert(
        ids=[f"{pdf.stem}-{i}" for i in range(len(chunks))],
        documents=chunks,
        metadatas=[{"sumber": pdf.name, "chunk": i} for i in range(len(chunks))],
    )
    print(f"{pdf.name} → {len(chunks)} chunk disimpan.")

print("\nSelesai! Indeks tersimpan di folder 'chroma_db'.")