"""
Lab P-10: RAG Part 1 - Embedding, Chunking, Vector DB
Cara pakai:
    pip install pypdf chromadb sentence-transformers
    mkdir dokumen   # taruh 2-3 PDF di sini

    python rag_lab.py lab1                 # Lab 1: eksperimen chunking
    python rag_lab.py lab2                 # Lab 2: indeks + 5 pertanyaan + 2 parafrase
    python rag_lab.py lab3                 # Lab 3: hit@k pada 2 ukuran chunk
    python rag_lab.py query "pertanyaan"   # tanya bebas (opsi --size 500)
"""
import argparse
import glob
import os
import random
import re

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

FOLDER = "dokumen"
DB_PATH = "chroma_db"
# Model multibahasa supaya bahasa Indonesia & parafrase dipahami dengan baik
EMBED_MODEL = "paraphrase-multilingual-MiniLM-L12-v2"

_ef = None
_client = None


# ---------------------------------------------------------------- dasar
def baca_pdf(path):
    """Baca seluruh teks PDF, rapikan spasi/baris baru."""
    reader = PdfReader(path)
    teks = " ".join((hal.extract_text() or "") for hal in reader.pages)
    return re.sub(r"\s+", " ", teks).strip()


def potong(teks, ukuran=500, overlap_persen=0.15):
    """Potong teks per `ukuran` karakter dengan overlap (10-20%)."""
    overlap = int(ukuran * overlap_persen)
    langkah = ukuran - overlap
    return [teks[i:i + ukuran] for i in range(0, len(teks), langkah)
            if teks[i:i + ukuran].strip()]


def get_client():
    global _client, _ef
    if _client is None:
        _client = chromadb.PersistentClient(path=DB_PATH)
        _ef = embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name=EMBED_MODEL)
    return _client


def daftar_pdf():
    files = sorted(glob.glob(os.path.join(FOLDER, "*.pdf")))
    if not files:
        raise SystemExit(f"Tidak ada PDF di folder '{FOLDER}/'. Letakkan 2-3 PDF dulu.")
    return files


# ---------------------------------------------------------------- indeks
def indeks(ukuran, overlap_persen=0.15, verbose=True):
    """Bangun koleksi `dok_<ukuran>` dari semua PDF di folder dokumen/."""
    client = get_client()
    nama = f"dok_{ukuran}"
    try:
        client.delete_collection(nama)
    except Exception:
        pass
    col = client.create_collection(nama, embedding_function=_ef,
                                   metadata={"hnsw:space": "cosine"})
    for path in daftar_pdf():
        chunks = potong(baca_pdf(path), ukuran, overlap_persen)
        fname = os.path.basename(path)
        col.add(
            ids=[f"{fname}-{i}" for i in range(len(chunks))],
            documents=chunks,
            metadatas=[{"sumber": fname, "chunk": i} for i in range(len(chunks))],
        )
        if verbose:
            print(f"  [{nama}] {fname}: {len(chunks)} chunk")
    return col


def get_collection(ukuran):
    client = get_client()
    try:
        return client.get_collection(f"dok_{ukuran}", embedding_function=_ef)
    except Exception:
        return indeks(ukuran, verbose=False)


def query(pertanyaan, k=3, ukuran=500, tampil=True):
    col = get_collection(ukuran)
    hasil = col.query(query_texts=[pertanyaan], n_results=k)
    docs = hasil["documents"][0]
    metas = hasil["metadatas"][0]
    jarak = hasil["distances"][0]
    if tampil:
        print(f"\nTanya: {pertanyaan}")
        for rank, (d, m, j) in enumerate(zip(docs, metas, jarak), 1):
            print(f"  #{rank} [{m['sumber']} chunk {m['chunk']}] jarak={j:.3f}")
            print(f"     {d[:200]}...")
    return docs


# ---------------------------------------------------------------- lab 1
def lab1():
    path = daftar_pdf()[0]
    teks = baca_pdf(path)
    print(f"=== LAB 1 | file: {path} | total {len(teks)} karakter ===")
    print("\n[1] 500 karakter pertama:\n")
    print(teks[:500])

    # --- Tabel 1: ringkasan jumlah chunk
    print("\n[2] Jumlah chunk (overlap 15%)\n")
    hasil = {}
    print(f"{'Ukuran':>7} | {'Overlap':>7} | {'Jml chunk':>9} | {'Rata2 panjang':>13} | {'Berakhir titik':>14}")
    print("-" * 64)
    for ukuran in (500, 1000, 1500):
        chunks = potong(teks, ukuran, 0.15)
        hasil[ukuran] = chunks
        rata = sum(len(c) for c in chunks) / len(chunks)
        titik = sum(c.rstrip().endswith(".") for c in chunks)
        print(f"{ukuran:>7} | {int(ukuran*0.15):>7} | {len(chunks):>9} | "
              f"{rata:>13.0f} | {titik:>8}/{len(chunks):<5}")

    # --- Tabel 2: 3 chunk acak per ukuran
    print("\n[3] 3 chunk acak per ukuran (cek: terpotong di tengah kalimat/tabel?)\n")
    print(f"{'Ukuran':>7} | {'Awal chunk (30 krk)':<32} | {'Akhir chunk (30 krk)':<32} | Titik?")
    print("-" * 85)
    for ukuran, chunks in hasil.items():
        for c in random.sample(chunks, min(3, len(chunks))):
            awal = c[:30].ljust(30)
            akhir = c.rstrip()[-30:].rjust(30)
            ya = "Ya" if c.rstrip().endswith(".") else "Tidak"
            print(f"{ukuran:>7} | {awal:<32} | ..{akhir:<30} | {ya}")
        print("-" * 85)


# ---------------------------------------------------------------- lab 2
PERTANYAAN_LAB2 = [
    "Wahana apa yang pertama kali mengirim gambar jarak dekat Mars?",
    "Siapa yang mengusulkan Tes Turing?",
    "Berapa target porsi energi terbarukan Indonesia tahun 2025?",
    "Apa itu Narrow AI dan contohnya?",
    "Apa saja tantangan pengembangan energi terbarukan?",
]
PERTANYAAN_PARAFRASE = [
    # tidak memakai kata yang sama dengan dokumen
    "Pesawat robotik pertama yang memotret planet merah dari dekat",
    "Sumber listrik dari perut bumi yang jalan nonstop tanpa peduli cuaca",
]


def lab2():
    print("=== LAB 2 | Bangun indeks (ukuran 500) ===")
    indeks(500)
    print("\n--- 5 pertanyaan: periksa manual apakah chunk teratas relevan ---")
    for q in PERTANYAAN_LAB2:
        query(q, k=3, ukuran=500)
    print("\n--- 2 pertanyaan parafrase: apakah pencarian semantik tetap menemukan? ---")
    for q in PERTANYAAN_PARAFRASE:
        query(q, k=3, ukuran=500)


# ---------------------------------------------------------------- lab 3
# (pertanyaan, [kata kunci jawaban]) -> hit jika SALAH SATU kata kunci ada di chunk
# Sesuaikan jika dokumen Anda berbeda.
SOAL_UJI = [
    ("Wahana apa yang pertama mengirim gambar jarak dekat Mars?", ["mariner 4"]),
    ("Kapan Viking 1 dan 2 mendarat di Mars?", ["1976"]),
    ("Apa nama rover pertama yang beroda di Mars?", ["sojourner"]),
    ("Apa temuan utama Curiosity?", ["bahan organik", "kawah gale"]),
    ("Apa program Mars Sample Return?", ["mars sample return", "msr"]),
    ("Siapa yang mengusulkan Tes Turing?", ["alan turing"]),
    ("Siapa yang mencetuskan istilah Artificial Intelligence pada 1956?", ["john mccarthy"]),
    ("Sebutkan contoh Narrow AI.", ["siri", "alexa"]),
    ("Siapa yang dikalahkan IBM Deep Blue?", ["kasparov"]),
    ("Siapa yang dikalahkan AlphaGo?", ["lee sedol"]),
    ("Berapa target bauran EBT nasional tahun 2025?", ["23%"]),
    ("Berapa persen cadangan panas bumi dunia yang dimiliki Indonesia?", ["40%"]),
    ("Berapa potensi tenaga surya Indonesia?", ["207,8"]),
    ("Berapa potensi tenaga air (hidro) Indonesia?", ["75,0"]),
    ("Apa tantangan pengembangan energi terbarukan di Indonesia?",
     ["investasi awal", "intermitensi", "transmisi"]),
]


def hit_at_k(ukuran, k):
    col = get_collection(ukuran)
    hit, gagal = 0, []
    for q, kunci in SOAL_UJI:
        docs = col.query(query_texts=[q], n_results=k)["documents"][0]
        gabung = " ".join(docs).lower()
        if any(kw in gabung for kw in kunci):
            hit += 1
        else:
            gagal.append(q)
    return hit / len(SOAL_UJI), gagal


def lab3():
    ukuran_list = (500, 1000)  # ubah ke (500, 1500) jika mau
    print("=== LAB 3 | hit@k pada dua ukuran chunk ===")
    for u in ukuran_list:
        indeks(u, verbose=False)

    print(f"\n{'ukuran':>8} | " + " | ".join(f"hit@{k}" for k in (1, 3, 5)))
    print("-" * 36)
    gagal_semua = {}
    for u in ukuran_list:
        baris = []
        for k in (1, 3, 5):
            skor, gagal = hit_at_k(u, k)
            baris.append(f"{skor:.2f}")
            if k == 1:
                gagal_semua[u] = gagal
        print(f"{u:>8} | " + " | ".join(f"{b:>5}" for b in baris))

    print("\nPertanyaan yang GAGAL pada k=1:")
    for u, g in gagal_semua.items():
        print(f"  ukuran {u}:")
        for q in g or ["(tidak ada)"]:
            print(f"    - {q}")
    print("\nTugas Anda: tulis kesimpulan 3-4 kalimat (ukuran mana lebih baik, "
          "pertanyaan apa gagal, dugaan penyebab).")


# ---------------------------------------------------------------- main
if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("perintah", choices=["lab1", "lab2", "lab3", "query"])
    ap.add_argument("teks", nargs="?", help="pertanyaan (untuk perintah query)")
    ap.add_argument("--size", type=int, default=500)
    ap.add_argument("-k", type=int, default=3)
    a = ap.parse_args()

    if a.perintah == "lab1":
        lab1()
    elif a.perintah == "lab2":
        lab2()
    elif a.perintah == "lab3":
        lab3()
    else:
        if not a.teks:
            raise SystemExit('Contoh: python rag_lab.py query "Siapa Alan Turing?"')
        query(a.teks, k=a.k, ukuran=a.size)