import os
from pathlib import Path

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


# =====================================================
# KONFIGURASI PROGRAM
# =====================================================

BASE_DIR = Path(__file__).resolve().parent
DB_PATH = str(BASE_DIR / "chroma_db")

UKURAN_CHUNK = [500, 1000, 1500]
OVERLAP_PERSEN = 0.10

MODEL_EMBEDDING = "paraphrase-multilingual-MiniLM-L12-v2"

print("Memuat model embedding...")
embedding = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name=MODEL_EMBEDDING
)

client = chromadb.PersistentClient(path=DB_PATH)


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

def ambil_file_pdf():
    return list(BASE_DIR.glob("*.pdf"))


def baca_pdf(path):
    reader = PdfReader(str(path))
    teks = ""

    for halaman in reader.pages:
        teks += (halaman.extract_text() or "") + "\n"

    return teks


def pilih_pdf(files):
    """Memilih satu PDF agar eksperimen konsisten."""
    if not files:
        print("Belum ada file PDF!")
        print("Letakkan PDF di folder yang sama dengan file Python ini.")
        return None

    if len(files) == 1:
        return files[0]

    print("\nDaftar dokumen PDF:")
    for i, file in enumerate(files, 1):
        print(f"{i}. {file.name}")

    while True:
        pilihan = input("Pilih nomor PDF: ").strip()
        if pilihan.isdigit() and 1 <= int(pilihan) <= len(files):
            return files[int(pilihan) - 1]
        print("Pilihan tidak valid.")


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

def chunking(teks, ukuran=1000, overlap=100):
    if ukuran <= 0 or overlap < 0 or overlap >= ukuran:
        raise ValueError("Ukuran atau overlap tidak valid.")

    chunks = []
    start = 0

    while start < len(teks):
        akhir = min(start + ukuran, len(teks))
        potongan = teks[start:akhir]

        if potongan.strip():
            chunks.append(potongan)

        start += ukuran - overlap

    return chunks


# =====================================================
# LAB 1 - EKSPERIMEN CHUNKING
# =====================================================

def lab_chunking():
    print("\n" + "=" * 60)
    print("LAB 1 - EKSPERIMEN CHUNKING")
    print("=" * 60)

    file = pilih_pdf(ambil_file_pdf())
    if file is None:
        return

    teks = baca_pdf(file)
    print(f"\nDokumen       : {file.name}")
    print(f"Panjang teks  : {len(teks)} karakter")

    if not teks.strip():
        print("Dokumen tidak memiliki teks yang dapat dibaca.")
        return

    for ukuran in UKURAN_CHUNK:
        for persentase in [0.10, 0.20]:
            overlap = int(ukuran * persentase)
            chunks = chunking(teks, ukuran, overlap)

            print("\n" + "-" * 55)
            print(f"Ukuran chunk : {ukuran}")
            print(f"Overlap      : {overlap} ({persentase:.0%})")
            print(f"Jumlah chunk : {len(chunks)}")
            print("Tiga contoh chunk:")

            # Tampilkan maksimal tiga contoh chunk.
            # Pemilihan posisi dibuat tersebar agar tidak selalu awal dokumen.
            if len(chunks) <= 3:
                indeks = list(range(len(chunks)))
            else:
                indeks = [0, len(chunks) // 2, len(chunks) - 1]

            for nomor, idx in enumerate(indeks, 1):
                print(f"\nContoh {nomor} (Chunk {idx + 1}):")
                print(chunks[idx][:350])

    print("\nEksperimen chunking selesai!")


# =====================================================
# FUNGSI MEMBANGUN KOLEKSI
# =====================================================

def bangun_koleksi(nama_koleksi, ukuran_chunk, file_pdf, reset=True):
    """Membangun koleksi ChromaDB dari satu dokumen PDF."""
    teks = baca_pdf(file_pdf)

    if not teks.strip():
        print("PDF kosong atau tidak dapat dibaca.")
        return None

    overlap = int(ukuran_chunk * OVERLAP_PERSEN)
    chunks = chunking(teks, ukuran_chunk, overlap)

    if reset:
        try:
            client.delete_collection(nama_koleksi)
        except Exception:
            pass

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

    ids = [f"chunk_{i}" for i in range(len(chunks))]
    metadatas = [
        {
            "sumber": file_pdf.name,
            "nomor_chunk": i,
            "ukuran_chunk": ukuran_chunk
        }
        for i in range(len(chunks))
    ]

    batch_size = 100
    for start in range(0, len(chunks), batch_size):
        end = start + batch_size
        collection.upsert(
            ids=ids[start:end],
            documents=chunks[start:end],
            metadatas=metadatas[start:end]
        )

    return collection


# =====================================================
# LAB 2 - MEMBANGUN INDEKS CHROMADB
# =====================================================

def lab_indexing():
    print("\n" + "=" * 60)
    print("LAB 2 - INDEXING CHROMADB")
    print("=" * 60)

    file = pilih_pdf(ambil_file_pdf())
    if file is None:
        return

    print(f"\nMembaca: {file.name}")
    collection = bangun_koleksi(
        nama_koleksi="rag_part1",
        ukuran_chunk=1000,
        file_pdf=file,
        reset=True
    )

    if collection is None:
        return

    print("\nIndexing berhasil!")
    print("Ukuran chunk :", 1000)
    print("Total chunk  :", collection.count())
    print("Database     :", DB_PATH)


# =====================================================
# MENGAMBIL KOLEKSI CHROMADB
# =====================================================

def ambil_koleksi():
    try:
        return client.get_collection(
            name="rag_part1",
            embedding_function=embedding
        )
    except Exception:
        print("\nDatabase belum tersedia!")
        print("Silakan jalankan Lab 2 terlebih dahulu.")
        return None


# =====================================================
# FUNGSI PENCARIAN DOKUMEN
# =====================================================

def cari_dokumen(pertanyaan, k=3, tampilkan=True, collection=None):
    if collection is None:
        collection = ambil_koleksi()

    if collection is None:
        return []

    total = collection.count()
    if total == 0:
        print("Database masih kosong!")
        return []

    k = min(k, total)

    hasil = collection.query(
        query_texts=[pertanyaan],
        n_results=k
    )

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

    if tampilkan:
        print("\nPertanyaan:", pertanyaan)
        print("\nHasil pencarian:")

        for i, doc in enumerate(dokumen):
            print("\n" + "-" * 45)
            print(f"HASIL KE-{i + 1}")
            print("-" * 45)
            print("Sumber:", metadata[i].get("sumber", "-"))
            print("Chunk:", metadata[i].get("nomor_chunk", "-"))
            print("Jarak:", round(jarak[i], 4))
            print("Isi:")
            print(doc[:500])

    return dokumen


# =====================================================
# LAB 2 - PENCARIAN DOKUMEN
# =====================================================

def lab_pencarian():
    print("\n" + "=" * 60)
    print("LAB 2 - PENCARIAN DOKUMEN")
    print("=" * 60)

    collection = ambil_koleksi()
    if collection is None or collection.count() == 0:
        print("Silakan lakukan indexing terlebih dahulu.")
        return

    print("\nMasukkan pertanyaan untuk mencari informasi.")
    print("Ketik '0' untuk kembali ke menu.")

    while True:
        pertanyaan = input("\nPertanyaan: ").strip()

        if pertanyaan == "0":
            break

        if not pertanyaan:
            print("Pertanyaan tidak boleh kosong.")
            continue

        cari_dokumen(pertanyaan, k=3, collection=collection)


# =====================================================
# LAB 2 - PENGUJIAN PARAFRASE
# =====================================================

def lab_parafrase():
    print("\n" + "=" * 60)
    print("LAB 2 - PENGUJIAN PARAFRASE")
    print("=" * 60)

    collection = ambil_koleksi()
    if collection is None or collection.count() == 0:
        print("Silakan lakukan indexing terlebih dahulu.")
        return

    pertanyaan_asli = input("\nPertanyaan asli: ").strip()
    if not pertanyaan_asli:
        print("Pertanyaan tidak boleh kosong.")
        return

    pertanyaan_parafrase = input("Pertanyaan parafrase: ").strip()
    if not pertanyaan_parafrase:
        print("Parafrase tidak boleh kosong.")
        return

    print("\n--- HASIL PERTANYAAN ASLI ---")
    cari_dokumen(pertanyaan_asli, k=3, collection=collection)

    print("\n--- HASIL PERTANYAAN PARAFRASE ---")
    cari_dokumen(pertanyaan_parafrase, k=3, collection=collection)


# =====================================================
# LAB 3 - EVALUASI HIT@K
# =====================================================

def buat_koleksi_evaluasi(ukuran_chunk, file_pdf):
    nama_koleksi = f"rag_eval_{ukuran_chunk}"

    print(f"\nMembangun koleksi {nama_koleksi}...")

    collection = bangun_koleksi(
        nama_koleksi=nama_koleksi,
        ukuran_chunk=ukuran_chunk,
        file_pdf=file_pdf,
        reset=True
    )

    if collection is not None:
        print("Indexing selesai.")
        print("Jumlah chunk:", collection.count())

    return collection


def hitung_hit_at_k(collection, data_uji, k):
    jumlah_hit = 0
    total = len(data_uji)

    for nomor, (pertanyaan, keyword) in enumerate(data_uji, 1):
        jumlah_dokumen = min(k, collection.count())

        hasil = collection.query(
            query_texts=[pertanyaan],
            n_results=jumlah_dokumen
        )

        dokumen = hasil["documents"][0]

        # Hit jika minimal satu dari K dokumen teratas
        # memuat kata/frasa kunci jawaban.
        ditemukan = any(
            keyword.casefold() in teks.casefold()
            for teks in dokumen
        )

        if ditemukan:
            jumlah_hit += 1

        status = "HIT" if ditemukan else "MISS"
        print(f"{nomor:02d}. {status} | {pertanyaan}")

    skor = jumlah_hit / total if total else 0
    return jumlah_hit, total, skor


def lab_evaluasi():
    print("\n" + "=" * 60)
    print("LAB 3 - EVALUASI HIT@K")
    print("=" * 60)

    print("""
Evaluasi menggunakan:
- 15 pertanyaan uji
- Dua koleksi: chunk 500 dan 1000 karakter
- Overlap 10%
- Hit@1, Hit@3, dan Hit@5

Masukkan kata/frasa yang benar-benar terdapat
pada bagian dokumen yang menjawab pertanyaan.
""")

    file_pdf = pilih_pdf(ambil_file_pdf())
    if file_pdf is None:
        return

    data_uji = []

    for i in range(15):
        print("\n" + "-" * 50)
        print(f"PERTANYAAN KE-{i + 1}")
        pertanyaan = input("Pertanyaan: ").strip()
        keyword = input("Kata kunci jawaban: ").strip()

        if not pertanyaan or not keyword:
            print("Pertanyaan dan kata kunci wajib diisi.")
            print("Evaluasi dibatalkan. Jalankan kembali menu 5.")
            return

        data_uji.append((pertanyaan, keyword))

    # Bangun dua koleksi dari PDF yang sama.
    koleksi_500 = buat_koleksi_evaluasi(500, file_pdf)
    if koleksi_500 is None:
        return

    koleksi_1000 = buat_koleksi_evaluasi(1000, file_pdf)
    if koleksi_1000 is None:
        return

    koleksi = {
        "Chunk 500": koleksi_500,
        "Chunk 1000": koleksi_1000
    }

    hasil_evaluasi = {}

    for nama, collection in koleksi.items():
        print("\n" + "=" * 60)
        print(f"EVALUASI {nama}")
        print("=" * 60)

        hasil_evaluasi[nama] = {}

        for k in [1, 3, 5]:
            print(f"\n--- PENGUJIAN HIT@{k} ---")

            hit, total, skor = hitung_hit_at_k(
                collection,
                data_uji,
                k
            )

            hasil_evaluasi[nama][k] = {
                "hit": hit,
                "total": total,
                "skor": skor
            }

            print(f"\nHit@{k} = {hit}/{total} = {skor:.2%}")

    # Tampilkan perbandingan.
    print("\n" + "=" * 68)
    print("TABEL PERBANDINGAN HASIL HIT@K")
    print("=" * 68)

    print(
        f"{'Ukuran Chunk':<18}"
        f"{'Hit@1':<16}"
        f"{'Hit@3':<16}"
        f"{'Hit@5':<16}"
    )
    print("-" * 68)

    for nama in koleksi:
        h1 = hasil_evaluasi[nama][1]["skor"]
        h3 = hasil_evaluasi[nama][3]["skor"]
        h5 = hasil_evaluasi[nama][5]["skor"]

        print(
            f"{nama:<18}"
            f"{h1:<16.2%}"
            f"{h3:<16.2%}"
            f"{h5:<16.2%}"
        )

    print("-" * 68)
    print("""
Catatan:
Hit@K adalah proporsi pertanyaan yang memiliki
minimal satu dokumen relevan pada K hasil teratas.

Evaluasi ini menggunakan pencocokan kata/frasa
kunci sederhana, bukan penilaian relevansi manual.
Periksa kembali hasil MISS dan kecocokan keyword
agar kesimpulan praktikum lebih akurat.
""")

    print("Evaluasi Hit@K selesai!")


# =====================================================
# MENU UTAMA
# =====================================================

def main():
    while True:
        print("\n")
        print("=" * 60)
        print("                 RAG PART 1 SYSTEM")
        print("=" * 60)
        print("1. Lab 1 - Eksperimen Chunking")
        print("2. Lab 2 - Indexing ChromaDB")
        print("3. Lab 2 - Pencarian Dokumen")
        print("4. Lab 2 - Pengujian Parafrase")
        print("5. Lab 3 - Evaluasi Hit@K")
        print("6. Keluar")
        print("=" * 60)

        pilihan = input("Pilih menu (1-6): ").strip()

        if pilihan == "1":
            lab_chunking()
        elif pilihan == "2":
            lab_indexing()
        elif pilihan == "3":
            lab_pencarian()
        elif pilihan == "4":
            lab_parafrase()
        elif pilihan == "5":
            lab_evaluasi()
        elif pilihan == "6":
            print("\nProgram selesai. Terima kasih!")
            break
        else:
            print("\nPilihan tidak valid!")


# =====================================================
# MENJALANKAN PROGRAM
# =====================================================

if __name__ == "__main__":
    main()