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


# ==========================================
# PENGATURAN
# ==========================================

FOLDER_DOKUMEN = Path("dokumen")
FOLDER_DATABASE = "./chroma_db"

UKURAN_CHUNK = [500, 1000, 1500]
MODEL_EMBEDDING = "paraphrase-multilingual-MiniLM-L12-v2"


# ==========================================
# MEMBACA FILE PDF
# ==========================================

def baca_pdf(path):
    reader = PdfReader(path)
    daftar_halaman = []

    for nomor, halaman in enumerate(reader.pages, start=1):
        teks = halaman.extract_text() or ""

        if teks.strip():
            daftar_halaman.append({
                "teks": teks,
                "halaman": nomor
            })

    return daftar_halaman


# ==========================================
# MEMOTONG TEKS MENJADI CHUNK
# ==========================================

def potong(teks, ukuran=500, tumpang=50):
    hasil = []
    mulai = 0

    while mulai < len(teks):
        bagian = teks[mulai:mulai + ukuran]

        if bagian.strip():
            hasil.append(bagian)

        mulai += ukuran - tumpang

    return hasil


# ==========================================
# MEMBUAT INDEKS CHROMADB
# ==========================================

def buat_indeks():
    if not FOLDER_DOKUMEN.exists():
        print("Folder dokumen tidak ditemukan.")
        return

    daftar_pdf = list(FOLDER_DOKUMEN.glob("*.pdf"))

    if not daftar_pdf:
        print("Belum ada file PDF di folder dokumen.")
        return

    print("Menghubungkan ke ChromaDB...")

    client = chromadb.PersistentClient(
        path=FOLDER_DATABASE
    )

    embedding = (
        embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name=MODEL_EMBEDDING
        )
    )

    print("Model embedding siap digunakan.")

    # Membuat koleksi untuk setiap ukuran chunk.
    for ukuran in UKURAN_CHUNK:

        nama_koleksi = f"dokumen_chunk_{ukuran}"

        koleksi = client.get_or_create_collection(
            name=nama_koleksi,
            embedding_function=embedding,
            metadata={"hnsw:space": "cosine"}
        )

        print("\n================================")
        print(f"MEMBUAT INDEKS CHUNK {ukuran}")
        print("================================")

        total_chunk = 0

        for pdf in daftar_pdf:
            halaman_pdf = baca_pdf(pdf)

            ids = []
            documents = []
            metadatas = []

            for data in halaman_pdf:
                teks = data["teks"]

                overlap = int(ukuran * 0.10)

                chunks = potong(
                    teks,
                    ukuran=ukuran,
                    tumpang=overlap
                )

                for nomor, chunk in enumerate(chunks):
                    id_chunk = (
                        f"{pdf.stem}_halaman_"
                        f"{data['halaman']}_chunk_"
                        f"{nomor}_ukuran_{ukuran}"
                    )

                    ids.append(id_chunk)
                    documents.append(chunk)

                    metadatas.append({
                        "sumber": pdf.name,
                        "halaman": data["halaman"],
                        "chunk": nomor,
                        "ukuran_chunk": ukuran
                    })

            if ids:
                # Simpan dalam kelompok agar lebih aman
                # untuk dokumen dengan banyak chunk.
                for awal in range(0, len(ids), 100):
                    akhir = awal + 100

                    koleksi.upsert(
                        ids=ids[awal:akhir],
                        documents=documents[awal:akhir],
                        metadatas=metadatas[awal:akhir]
                    )

                print(
                    f"{pdf.name}: {len(ids)} chunk "
                    f"untuk ukuran {ukuran}"
                )

                total_chunk += len(ids)

            else:
                print(
                    f"{pdf.name}: tidak ada teks yang terbaca."
                )

        print(f"Total chunk: {total_chunk}")

    print("\nSemua indeks berhasil dibuat!")
    print(f"Lokasi database: {FOLDER_DATABASE}")


# ==========================================
# PROGRAM UTAMA
# ==========================================

if __name__ == "__main__":
    buat_indeks()