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

# ==========================================
# PROGRAM RAG P-10
# EMBEDDING & VECTOR DATABASE
# ==========================================

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

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

# Membuat model embedding dan database
print("Memuat model embedding...")
embedding = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name=MODEL_EMBEDDING
)

client = chromadb.PersistentClient(path=FOLDER_DATABASE)


# ==========================================
# 1. MEMBACA PDF
# ==========================================

def baca_pdf(path):
    reader = PdfReader(str(path))
    hasil = []

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

        if teks.strip():
            hasil.append((nomor, teks))

    return hasil


# ==========================================
# 2. CHUNKING
# ==========================================

def potong_teks(teks, ukuran, overlap=OVERLAP):
    chunks = []
    langkah = ukuran - overlap

    if langkah <= 0:
        raise ValueError("Overlap harus lebih kecil dari ukuran chunk.")

    mulai = 0

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

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

        mulai += langkah

    return chunks


# ==========================================
# 3. MEMBANGUN INDEKS CHROMADB
# ==========================================

def bangun_indeks():
    if not FOLDER_DOKUMEN.exists():
        FOLDER_DOKUMEN.mkdir(parents=True)
        print("Folder dokumen berhasil dibuat.")
        print("Masukkan file PDF ke folder dokumen, lalu jalankan lagi.")
        return

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

    if not daftar_pdf:
        print("Tidak ada PDF di folder dokumen.")
        print("Masukkan minimal satu PDF berbasis teks.")
        return

    for ukuran in UKURAN_CHUNK:
        nama_koleksi = f"pedoman_{ukuran}"

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

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

        for path_pdf in daftar_pdf:
            try:
                halaman_pdf = baca_pdf(path_pdf)

                semua_id = []
                semua_teks = []
                semua_metadata = []

                for nomor_halaman, teks in halaman_pdf:
                    chunks = potong_teks(
                        teks,
                        ukuran,
                        min(OVERLAP, ukuran - 1)
                    )

                    for nomor_chunk, chunk in enumerate(chunks):
                        identitas = (
                            f"{path_pdf.stem}_"
                            f"{nomor_halaman}_"
                            f"{nomor_chunk}"
                        )

                        semua_id.append(
                            f"{identitas}_{ukuran}"
                        )
                        semua_teks.append(chunk)

                        semua_metadata.append({
                            "sumber": path_pdf.name,
                            "halaman": nomor_halaman,
                            "chunk": nomor_chunk,
                            "ukuran_chunk": ukuran
                        })

                # Menghindari data ganda saat indexing ulang
                data_lama = koleksi.get(
                    where={"sumber": path_pdf.name}
                )

                if data_lama["ids"]:
                    koleksi.delete(ids=data_lama["ids"])

                if semua_teks:
                    koleksi.upsert(
                        ids=semua_id,
                        documents=semua_teks,
                        metadatas=semua_metadata
                    )

                print(
                    f"{path_pdf.name}: "
                    f"{len(semua_teks)} chunk"
                )

            except Exception as error:
                print(f"Gagal memproses {path_pdf.name}: {error}")

    print("\nIndexing selesai.")


# ==========================================
# 4. PENCARIAN SEMANTIK
# ==========================================

def cari_dokumen():
    pertanyaan = input("\nMasukkan pertanyaan: ").strip()

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

    try:
        ukuran = int(input("Ukuran chunk (500/1000/1500): "))

        if ukuran not in UKURAN_CHUNK:
            print("Ukuran chunk tidak valid.")
            return

        nama_koleksi = f"pedoman_{ukuran}"
        koleksi = client.get_collection(
            name=nama_koleksi,
            embedding_function=embedding
        )

        jumlah_data = koleksi.count()

        if jumlah_data == 0:
            print("Koleksi kosong. Jalankan indexing dahulu.")
            return

        jumlah = min(3, jumlah_data)

        hasil = koleksi.query(
            query_texts=[pertanyaan],
            n_results=jumlah
        )

        print("\n=== HASIL PENCARIAN ===")
        print("Pertanyaan:", pertanyaan)
        print("Ukuran chunk:", ukuran)

        for nomor, (teks, metadata, jarak) in enumerate(
            zip(
                hasil["documents"][0],
                hasil["metadatas"][0],
                hasil["distances"][0]
            ),
            start=1
        ):
            print(f"\nHasil ke-{nomor}")
            print("Sumber:", metadata["sumber"])
            print("Halaman:", metadata["halaman"])
            print("Chunk:", metadata["chunk"])
            print(f"Jarak cosine: {jarak:.4f}")
            print("Isi:", teks[:500])

    except ValueError:
        print("Masukkan angka ukuran chunk yang valid.")
    except Exception as error:
        print("Gagal melakukan pencarian:", error)


# ==========================================
# 5. EVALUASI HIT@K DAN MRR
# ==========================================

# PENTING:
# Ganti pertanyaan dan kata kunci berikut berdasarkan PDF milikmu.
# Kata kunci harus benar-benar muncul pada chunk yang menjawab pertanyaan.

UJI = [
    ("Apa syarat cuti akademik?", "cuti akademik"),
    ("Berapa batas kehadiran untuk ujian?", "75 persen"),
    ("Bagaimana prosedur pendaftaran?", "pendaftaran"),
    ("Apa persyaratan yang harus dipenuhi?", "persyaratan"),
    ("Berapa lama batas waktu yang ditentukan?", "batas waktu"),
    ("Siapa yang memberikan persetujuan?", "persetujuan"),
    ("Apa ketentuan bagi mahasiswa?", "mahasiswa"),
    ("Bagaimana cara mengajukan permohonan?", "permohonan"),
    ("Apa akibat pelanggaran aturan?", "aturan"),
    ("Kapan registrasi dilakukan?", "registrasi"),
    ("Apa dokumen yang harus disiapkan?", "dokumen"),
    ("Bagaimana proses evaluasi dilakukan?", "evaluasi"),
    ("Apa ketentuan mengenai biaya?", "biaya"),
    ("Siapa yang bertanggung jawab?", "bertanggung jawab"),
    ("Apa tujuan prosedur tersebut?", "tujuan")
]


def evaluasi_retrieval():
    print("\n=== EVALUASI RETRIEVAL ===")
    print("Pastikan pertanyaan dan kata kunci sesuai isi PDF.")

    for ukuran in UKURAN_CHUNK:
        nama_koleksi = f"pedoman_{ukuran}"

        try:
            koleksi = client.get_collection(
                name=nama_koleksi,
                embedding_function=embedding
            )

            jumlah_data = koleksi.count()

            if jumlah_data == 0:
                print(f"\nKoleksi {ukuran} kosong.")
                continue

            print(f"\n--- CHUNK {ukuran} ---")

            for k in [1, 3, 5]:
                jumlah_hit = 0
                total_reciprocal_rank = 0

                for pertanyaan, kata_kunci in UJI:
                    jumlah_hasil = min(k, jumlah_data)

                    hasil = koleksi.query(
                        query_texts=[pertanyaan],
                        n_results=jumlah_hasil
                    )

                    dokumen = hasil["documents"][0]
                    peringkat_benar = 0

                    for peringkat, teks in enumerate(
                        dokumen, start=1
                    ):
                        if kata_kunci.lower() in teks.lower():
                            peringkat_benar = peringkat
                            break

                    if peringkat_benar > 0:
                        jumlah_hit += 1
                        total_reciprocal_rank += (
                            1 / peringkat_benar
                        )

                hit_k = jumlah_hit / len(UJI)
                mrr = total_reciprocal_rank / len(UJI)

                print(
                    f"Hit@{k}: {hit_k:.2f} | "
                    f"MRR: {mrr:.2f}"
                )

        except Exception as error:
            print(f"Gagal mengevaluasi chunk {ukuran}: {error}")


# ==========================================
# 6. MENU UTAMA
# ==========================================

def menu():
    while True:
        print("\n" + "=" * 45)
        print("       PROGRAM RAG P-10")
        print(" Embedding & Vector Database")
        print("=" * 45)
        print("1. Membaca PDF dan membangun indeks")
        print("2. Melakukan pencarian semantik")
        print("3. Evaluasi Hit@k dan MRR")
        print("4. Keluar")
        print("=" * 45)

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

        if pilihan == "1":
            bangun_indeks()

        elif pilihan == "2":
            cari_dokumen()

        elif pilihan == "3":
            evaluasi_retrieval()

        elif pilihan == "4":
            print("Program selesai.")
            break

        else:
            print("Pilihan tidak valid. Masukkan angka 1-4.")


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

if __name__ == "__main__":
    menu()