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

# FUNGSI MEMBACA PDF
def baca_pdf(path):
    reader = PdfReader(path)

    return "\n".join(
        (p.extract_text() or "")
        for p in reader.pages
    )

# FUNGSI MEMOTONG TEKS
def potong(teks, ukuran=800, tumpang=100):

    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()]

# MENCARI PDF DI FOLDER DOKUMEN
daftar_pdf = list(Path("dokumen").glob("*.pdf"))

if not daftar_pdf:
    print("Tidak ada file PDF di folder dokumen!")

else:
    # Mengambil PDF pertama
    pdf = daftar_pdf[0]

    # Membaca isi PDF
    teks = baca_pdf(pdf)

    print(f"\nNama PDF: {pdf.name}")
    print(f"Total karakter: {len(teks)}")

    # Menampilkan 500 karakter pertama
    print("\n===== 500 KARAKTER PERTAMA =====")
    print(teks[:500])

    # Eksperimen tiga ukuran chunk
    for ukuran in (500, 1000, 1500):

        overlap = int(ukuran * 0.1)

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

        print("\n" + "=" * 50)
        print(f"UKURAN CHUNK: {ukuran}")
        print(f"OVERLAP: {overlap}")
        print(f"JUMLAH CHUNK: {len(chunks)}")

        # Mengambil maksimal 3 chunk secara acak
        sampel = random.sample(
            range(len(chunks)),
            min(3, len(chunks))
        )

        for nomor in sampel:
            print(f"\n--- CHUNK {nomor} ---")
            print(chunks[nomor])


# LAB 2: EMBEDDING DAN CHROMADB

if daftar_pdf:

    # Model embedding multilingual
    emb = embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name="paraphrase-multilingual-MiniLM-L12-v2"
    )

    # Membuat database lokal
    client = chromadb.PersistentClient(path="./chroma_db")

    def bangun_indeks(nama_koleksi, ukuran):

        # Menghapus koleksi lama agar tidak tercampur
        try:
            client.delete_collection(name=nama_koleksi)
        except Exception:
            pass

        # Membuat koleksi baru
        koleksi = client.create_collection(
            name=nama_koleksi,
            embedding_function=emb,
            metadata={"hnsw:space": "cosine"}
        )

        for pdf in daftar_pdf:

            teks = baca_pdf(pdf)

            overlap = int(ukuran * 0.1)

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

            if not chunks:
                print(pdf.name, "tidak memiliki teks.")
                continue

            koleksi.upsert(
                ids=[
                    f"{pdf.stem}-{ukuran}-{i}"
                    for i in range(len(chunks))
                ],
                documents=chunks,
                metadatas=[
                    {
                        "sumber": pdf.name,
                        "chunk": i,
                        "ukuran": ukuran
                    }
                    for i in range(len(chunks))
                ]
            )

            print(
                pdf.name,
                "->",
                len(chunks),
                "chunk"
            )

        print(
            f"\nKoleksi {nama_koleksi} berhasil dibuat."
        )
        print("Total data:", koleksi.count())


    # Membuat dua koleksi untuk eksperimen
    bangun_indeks("pedoman_500", 500)
    bangun_indeks("pedoman_1000", 1000)

else:
    print("Tidak ada PDF untuk indexing.")