"""
rag_lab.py - Kode untuk Lab 1, 2, dan 3 (chunking, indeks & pencarian, Hit@k)

Instalasi:
    pip install pypdf chromadb sentence-transformers

Cara pakai:
    python rag_lab.py lab1 dokumen/contoh.pdf
    python rag_lab.py lab2                      # indeks semua PDF di dokumen/ lalu tanya
    python rag_lab.py lab3 pertanyaan_uji.json  # hitung Hit@k

Format pertanyaan_uji.json:
[
  {"pertanyaan": "Berapa batas usia pendaftar?", "kata_kunci": ["usia", "35 tahun"]},
  ...
]
"""

import os

os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"  # sembunyikan peringatan symlink di Windows

import json
import random
import sys
from pathlib import Path

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

FOLDER_DOKUMEN = Path("dokumen")
FOLDER_DB = "chroma_db"
# Model multibahasa, cocok untuk teks Indonesia
MODEL_EMBEDDING = "paraphrase-multilingual-MiniLM-L12-v2"

_embed_fn = None
_client = None


def _embedding():
    global _embed_fn
    if _embed_fn is None:
        _embed_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name=MODEL_EMBEDDING
        )
    return _embed_fn


def _db():
    global _client
    if _client is None:
        _client = chromadb.PersistentClient(path=FOLDER_DB)
    return _client


# ---------------------------------------------------------------- Lab 1
def baca_pdf(path):
    """Baca seluruh teks PDF menjadi satu string."""
    reader = PdfReader(str(path))
    halaman = []
    for p in reader.pages:
        teks = p.extract_text() or ""
        halaman.append(teks.strip())
    return "\n\n".join(halaman)


def potong(teks, ukuran=1000, overlap_persen=15):
    """
    Potong teks menjadi chunk berukuran ~`ukuran` karakter dengan overlap.
    Berusaha memotong di akhir kalimat/paragraf agar tidak terpotong di tengah.
    """
    overlap = int(ukuran * overlap_persen / 100)
    chunks = []
    mulai = 0
    n = len(teks)

    while mulai < n:
        akhir = min(mulai + ukuran, n)

        if akhir < n:
            # cari titik potong terbaik di 25% terakhir jendela
            batas_cari = mulai + int(ukuran * 0.75)
            jendela = teks[batas_cari:akhir]
            titik = max(jendela.rfind("\n\n"), jendela.rfind(". "), jendela.rfind(".\n"))
            if titik != -1:
                akhir = batas_cari + titik + 1

        potongan = teks[mulai:akhir].strip()
        if potongan:
            chunks.append(potongan)

        if akhir >= n:
            break
        mulai_baru = max(akhir - overlap, mulai + 1)
        # geser maju sampai awal kata berikutnya agar tidak memotong kata
        while mulai_baru < akhir and not teks[mulai_baru - 1].isspace():
            mulai_baru += 1
        mulai = mulai_baru

    return chunks


def lab1(path_pdf):
    teks = baca_pdf(path_pdf)
    print(f"Total karakter: {len(teks)}")
    print("\n=== 500 karakter pertama ===")
    print(teks[:500])

    random.seed(42)
    for ukuran in (500, 1000, 1500):
        overlap_persen = 15
        chunks = potong(teks, ukuran, overlap_persen)
        print(f"\n{'=' * 60}")
        print(f"Ukuran {ukuran}, overlap {overlap_persen}% -> {len(chunks)} chunk")
        for i in random.sample(range(len(chunks)), min(3, len(chunks))):
            print(f"\n--- chunk #{i} ({len(chunks[i])} karakter) ---")
            print(chunks[i])
            akhir = chunks[i].rstrip()[-1:]
            if akhir not in ".!?:)":
                print("[!] Chunk tidak berakhir di tanda baca: kemungkinan terpotong")


# ---------------------------------------------------------------- Lab 2
def buat_indeks(nama_koleksi, ukuran, overlap_persen=15, folder=FOLDER_DOKUMEN):
    """Indeks semua PDF di `folder` ke koleksi baru bernama `nama_koleksi`."""
    try:
        _db().delete_collection(nama_koleksi)
    except Exception:
        pass
    koleksi = _db().create_collection(
        nama_koleksi,
        embedding_function=_embedding(),
        metadata={"hnsw:space": "cosine"},
    )

    pdfs = sorted(Path(folder).glob("*.pdf"))
    if not pdfs:
        sys.exit(f"Tidak ada PDF di folder '{folder}/'.")

    for pdf in pdfs:
        chunks = potong(baca_pdf(pdf), ukuran, overlap_persen)
        if not chunks:
            print(f"[!] {pdf.name}: tidak ada teks terbaca (PDF hasil scan?)")
            continue
        koleksi.add(
            ids=[f"{pdf.stem}-{i}" for i in range(len(chunks))],
            documents=chunks,
            metadatas=[{"sumber": pdf.name, "urutan": i} for i in range(len(chunks))],
        )
        print(f"{pdf.name}: {len(chunks)} chunk -> koleksi '{nama_koleksi}'")
    return koleksi


def ambil_koleksi(nama_koleksi):
    return _db().get_collection(nama_koleksi, embedding_function=_embedding())


def query(nama_koleksi, pertanyaan, k=3, tampil=True):
    """Cari k chunk paling mirip. Mengembalikan list dict."""
    hasil = ambil_koleksi(nama_koleksi).query(query_texts=[pertanyaan], n_results=k)
    out = []
    for doc, meta, jarak in zip(
        hasil["documents"][0], hasil["metadatas"][0], hasil["distances"][0]
    ):
        out.append({"teks": doc, "sumber": meta["sumber"], "skor": 1 - jarak})
    if tampil:
        print(f"\nPertanyaan: {pertanyaan}")
        for r in out:
            print(f"  [{r['skor']:.3f}] {r['sumber']}: {r['teks'][:700].replace(chr(10), ' ')}...\n")
    return out


def lab2():
    buat_indeks("chunk_1000", ukuran=1000)
    print("\nKetik pertanyaan (kosong untuk selesai). Coba 5 biasa + 2 parafrase.")
    while True:
        q = input("\n> ").strip()
        if not q:
            break
        query("chunk_1000", q, k=3)


# ---------------------------------------------------------------- Lab 3
def hit_at_k(nama_koleksi, soal, k):
    """
    Hit@k = proporsi pertanyaan yang kata kuncinya muncul
    (minimal satu kata kunci) di salah satu dari k chunk teratas.
    Mengembalikan (skor, daftar_pertanyaan_gagal).
    """
    hit = 0
    gagal = []
    for s in soal:
        hasil = query(nama_koleksi, s["pertanyaan"], k=k, tampil=False)
        gabungan = " ".join(" ".join(r["teks"].split()) for r in hasil).lower()
        if any(" ".join(kw.split()).lower() in gabungan for kw in s["kata_kunci"]):
            hit += 1
        else:
            gagal.append(s["pertanyaan"])
    return hit / len(soal), gagal


def lab3(path_soal):
    soal = json.loads(Path(path_soal).read_text(encoding="utf-8"))
    print(f"{len(soal)} pertanyaan uji dimuat.")

    ukuran_uji = {"chunk_500": 500, "chunk_1500": 1500}
    for nama, ukuran in ukuran_uji.items():
        buat_indeks(nama, ukuran)

    print(f"\n{'Koleksi':<12}{'Hit@1':>8}{'Hit@3':>8}{'Hit@5':>8}")
    gagal_per_koleksi = {}
    for nama in ukuran_uji:
        baris = []
        for k in (1, 3, 5):
            skor, gagal = hit_at_k(nama, soal, k)
            baris.append(skor)
            if k == 5:
                gagal_per_koleksi[nama] = gagal
        print(f"{nama:<12}" + "".join(f"{b:>8.2f}" for b in baris))

    print("\nPertanyaan yang gagal pada k=5:")
    for nama, gagal in gagal_per_koleksi.items():
        print(f"\n[{nama}] {len(gagal)} gagal")
        for g in gagal:
            print(f"  - {g}")

    print("\nTulis kesimpulan 3-4 kalimat: ukuran terbaik, jenis pertanyaan yang gagal, dugaan penyebab.")


# ---------------------------------------------------------------- CLI
if __name__ == "__main__":
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    perintah = sys.argv[1]
    if perintah == "lab1" and len(sys.argv) == 3:
        lab1(sys.argv[2])
    elif perintah == "lab2":
        lab2()
    elif perintah == "lab3" and len(sys.argv) == 3:
        lab3(sys.argv[2])
    else:
        sys.exit(__doc__)