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

# ==========================================
# LAB 2 - BANGUN INDEKS & CARI
# ==========================================

print("=" * 65)
print("LAB 2 - BANGUN INDEKS & CARI")
print("=" * 65, flush=True)

# ==========================================
# KONFIGURASI
# ==========================================

FOLDER_PDF = Path("dokumen")
DATABASE = "./chroma_db"

UKURAN_CHUNK = 800
OVERLAP = 100

# ==========================================
# FUNGSI MEMBACA PDF
# ==========================================

def baca_pdf(path):

    reader = PdfReader(path)

    halaman = []

    for nomor, page in enumerate(reader.pages, start=1):

        teks = page.extract_text() or ""

        halaman.append(teks)

    return "\n".join(halaman)


# ==========================================
# FUNGSI CHUNKING
# ==========================================

def potong(teks, ukuran=800, tumpang=100):

    hasil = []
    mulai = 0

    if ukuran <= 0 or tumpang < 0 or tumpang >= ukuran:
        raise ValueError("Pengaturan chunk tidak valid.")

    while mulai < len(teks):

        chunk = teks[mulai:mulai + ukuran]

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

        mulai += ukuran - tumpang

    return hasil


# ==========================================
# EMBEDDING MODEL
# ==========================================

print("\n[1] Memuat model embedding...", flush=True)

emb = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="paraphrase-multilingual-MiniLM-L12-v2",
    device="cpu"
)

print("Model embedding berhasil dimuat!", flush=True)


# ==========================================
# KONEKSI CHROMADB
# ==========================================

print("\n[2] Menghubungkan database ChromaDB...", flush=True)

client = chromadb.PersistentClient(path=DATABASE)

print("Database berhasil dihubungkan!", flush=True)


# ==========================================
# LANGKAH 1 - INDEXING PDF
# ==========================================

print("\n[3] Memeriksa file PDF...", flush=True)

daftar_pdf = sorted(FOLDER_PDF.glob("*.pdf"))

if not 2 <= len(daftar_pdf) <= 3:

    raise SystemExit(
        f"ERROR: Harus tersedia 2-3 PDF. "
        f"Saat ini ditemukan {len(daftar_pdf)} PDF."
    )

print(f"Ditemukan {len(daftar_pdf)} file PDF.", flush=True)

# Membaca dan memotong semua PDF
data_index = []

for pdf in daftar_pdf:

    print(f"\nMembaca: {pdf.name}", flush=True)

    teks = baca_pdf(pdf)

    if not teks.strip():
        print("PDF tidak memiliki teks yang bisa diekstrak.")
        continue

    chunks = potong(
        teks,
        ukuran=UKURAN_CHUNK,
        tumpang=OVERLAP
    )

    data_index.append((pdf, chunks))

    print(
        f"Berhasil membaca {len(chunks)} chunk.",
        flush=True
    )

if not data_index:
    raise SystemExit("Tidak ada PDF yang berhasil dibaca.")

# Menghapus koleksi lama agar tidak ada data duplikat
print("\nMembersihkan indeks lama...", flush=True)

try:
    client.delete_collection("pedoman")
except Exception:
    pass

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

print("Koleksi baru berhasil dibuat!", flush=True)


# ==========================================
# MEMASUKKAN DATA KE DATABASE
# ==========================================

print("\n=== LAPORAN INDEXING ===", flush=True)

total_chunk = 0

for pdf, chunks in data_index:

    ids = [
        f"{pdf.stem}-{i}"
        for i in range(len(chunks))
    ]

    metadatas = [
        {
            "sumber": pdf.name,
            "chunk": i
        }
        for i in range(len(chunks))
    ]

    print(
        f"Mengindeks {pdf.name}...",
        flush=True
    )

    # Embedding dilakukan secara bertahap
    batch_size = 32

    for mulai in range(0, len(chunks), batch_size):

        batch_chunks = chunks[mulai:mulai + batch_size]

        batch_ids = ids[mulai:mulai + batch_size]

        batch_meta = metadatas[mulai:mulai + batch_size]

        koleksi.upsert(
            ids=batch_ids,
            documents=batch_chunks,
            metadatas=batch_meta
        )

        print(
            f"  Proses: {min(mulai + batch_size, len(chunks))}"
            f"/{len(chunks)} chunk",
            flush=True
        )

    total_chunk += len(chunks)

    print(
        f"{pdf.name} -> {len(chunks)} chunk",
        flush=True
    )

print("\nTOTAL SELURUH CHUNK:", total_chunk, flush=True)

print(
    "Jumlah data di ChromaDB:",
    koleksi.count(),
    flush=True
)


# ==========================================
# FUNGSI PENCARIAN QUERY
# ==========================================

def tampilkan_query(q, label):

    print("\n" + "-" * 65)
    print(f"[{label}] {q}")
    print("-" * 65, flush=True)

    try:

        jumlah_data = koleksi.count()

        if jumlah_data == 0:
            print("Database kosong!")
            return

        # Tahap 1: Membuat embedding pertanyaan
        print(
            "  [1/3] Membuat embedding pertanyaan...",
            flush=True
        )

        vektor = emb([q])

        print(
            "  [2/3] Mencari chunk yang relevan...",
            flush=True
        )

        # Tahap 2: Pencarian menggunakan vektor
        hasil = koleksi.query(
            query_embeddings=vektor,
            n_results=min(3, jumlah_data),
            include=[
                "documents",
                "metadatas",
                "distances"
            ]
        )

        print(
            "  [3/3] Menampilkan hasil pencarian...",
            flush=True
        )

        dokumen = hasil["documents"][0]
        metadata = hasil["metadatas"][0]
        jarak = hasil["distances"][0]

        for nomor, (doc, meta, dist) in enumerate(
            zip(dokumen, metadata, jarak),
            start=1
        ):

            print(f"\n  HASIL {nomor}")

            print(f"  Sumber : {meta['sumber']}")
            print(f"  Chunk  : {meta['chunk']}")
            print(f"  Jarak  : {dist:.4f}")

            print("\n  Isi chunk:")
            print(" ", doc[:250].replace("\n", " "))

            # Evaluasi awal berdasarkan jarak cosine
            if dist < 0.3:
                print("  Kemiripan: Relatif tinggi")
            elif dist < 0.6:
                print("  Kemiripan: Sedang, perlu diperiksa")
            else:
                print("  Kemiripan: Relatif rendah")

        print("\nQuery berhasil diselesaikan!", flush=True)

    except Exception as error:

        print(
            f"\nERROR SAAT QUERY: {error}",
            flush=True
        )


# ==========================================
# LANGKAH 2 - 5 PERTANYAAN LANGSUNG
# ==========================================

pertanyaan = [

    "What is the basic concept of probability?",

    "What is the multiplication rule?",

    "What is the addition rule?",

    "What is a sample space?",

    "What is the difference between classical and empirical probability?"

]

print("\n" + "=" * 65)
print("LANGKAH 2 - LIMA PERTANYAAN LANGSUNG")
print("=" * 65, flush=True)

for nomor, q in enumerate(pertanyaan, start=1):

    print(
        f"\nMemproses pertanyaan {nomor}/5...",
        flush=True
    )

    tampilkan_query(q, "PERTANYAAN LANGSUNG")

    time.sleep(0.3)


# ==========================================
# LANGKAH 3 - 2 PERTANYAAN PARAFRASE
# ==========================================

parafrase = [

    "How do you find the chance of two events happening together?",

    "What is the collection of all possible outcomes in an experiment?"

]

print("\n" + "=" * 65)
print("LANGKAH 3 - DUA PERTANYAAN PARAFRASE")
print("=" * 65, flush=True)

for nomor, q in enumerate(parafrase, start=1):

    print(
        f"\nMemproses parafrase {nomor}/2...",
        flush=True
    )

    tampilkan_query(q, "PERTANYAAN PARAFRASE")

    time.sleep(0.3)


# ==========================================
# KESIMPULAN
# ==========================================

print("\n" + "=" * 65)
print("KESIMPULAN EKSPERIMEN LAB 2")
print("=" * 65)

print(f"Jumlah PDF       : {len(data_index)}")
print(f"Total chunk      : {total_chunk}")
print(f"Pertanyaan biasa : {len(pertanyaan)}")
print(f"Pertanyaan parafrase: {len(parafrase)}")

print("\nIndexing dan pencarian selesai!")
print("Silakan periksa relevansi setiap hasil pencarian.")
print("=" * 65, flush=True)