import os
import random

# Gunakan cache model lokal jika sudah terunduh agar tidak lambat saat offline
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

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

def baca_pdf(path):
    reader = PdfReader(path)
    return "\n".join((p.extract_text() or "") for p in reader.pages)

def potong(teks, ukuran=800, tumpang=100):
    """Pemotongan ukuran tetap dengan overlap (versi paling sederhana)."""
    hasil, mulai = [], 0
    while mulai < len(teks):
        hasil.append(teks[mulai:mulai + ukuran])
        mulai += ukuran - tumpang
    return [h for h in hasil if h.strip()]


# model embedding multilingual (jalan lokal, tanpa API key)
emb = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="paraphrase-multilingual-MiniLM-L12-v2")

client = chromadb.PersistentClient(path="./chroma_db")   # tersimpan di folder

# Koleksi utama (ukuran chunk default 800, overlap 100)
koleksi = client.get_or_create_collection(
    name="pedoman", embedding_function=emb,
    metadata={"hnsw:space": "cosine"})                    # ukuran jarak: cosine

# Dua koleksi tambahan untuk Lab 3 (perbandingan dua ukuran chunk berbeda)
koleksi_kecil = client.get_or_create_collection(
    name="pedoman_chunk_500", embedding_function=emb,
    metadata={"hnsw:space": "cosine"})

koleksi_besar = client.get_or_create_collection(
    name="pedoman_chunk_1500", embedding_function=emb,
    metadata={"hnsw:space": "cosine"})


def eksperimen_chunking(folder="dokumen"):
    """Lab 1 — Eksperimen chunking (ukuran 500, 1000, 1500 dengan overlap 15%)."""
    daftar_pdf = sorted(Path(folder).glob("*.pdf"))
    if not daftar_pdf:
        print(f"Tidak ada file PDF di '{folder}/'. Masukkan minimal 1 PDF terlebih dahulu.")
        return

    # Pilih PDF dengan teks terpanjang agar ukuran 1500 pun memiliki >= 3 chunk
    pdf_pertama = max(daftar_pdf, key=lambda p: len(baca_pdf(p)))
    teks = baca_pdf(pdf_pertama)

    print("=" * 70)
    print(f"LAB 1 — EKSPERIMEN CHUNKING (Sampel: {pdf_pertama.name})")
    print("=" * 70)
    print("\n[1] 500 Karakter Pertama:")
    print("-" * 70)
    print(teks.lstrip()[:500])
    print("-" * 70)

    # Tiga pengaturan ukuran chunk (500, 1000, 1500) dengan overlap 15% (antara 10-20%)
    pengaturan = [
        (500, 75),    # overlap 15%
        (1000, 150),  # overlap 15%
        (1500, 225),  # overlap 15%
    ]

    random.seed(42)
    print("\n[2] Perbandingan Jumlah Chunk & 3 Sampel Acak Tiap Ukuran:")
    for ukuran, tumpang in pengaturan:
        chunks = potong(teks, ukuran=ukuran, tumpang=tumpang)
        print(f"\n>>> Ukuran = {ukuran}, Overlap = {tumpang} ({tumpang*100//ukuran}%) → Total: {len(chunks)} chunk")
        sampel = random.sample(chunks, k=min(3, len(chunks)))
        for idx, c in enumerate(sampel, start=1):
            cuplikan = c.replace("\n", " ")
            print(f"  Sampel {idx}: \"{cuplikan[:120]} ... {cuplikan[-80:]}\"")


def indeks_ke_koleksi(target_koleksi, daftar_pdf, ukuran=800, tumpang=100, tampilkan_log=True):
    """Memotong semua PDF dan melakukan upsert ke koleksi ChromaDB tertentu."""
    for pdf in daftar_pdf:
        chunks = potong(baca_pdf(pdf), ukuran=ukuran, tumpang=tumpang)
        if not chunks:
            if tampilkan_log:
                print(f"{pdf.name} → dilewati (tidak ada teks)")
            continue
        target_koleksi.upsert(
            ids=[f"{pdf.stem}-{i}" for i in range(len(chunks))],
            documents=chunks,
            metadatas=[{"sumber": pdf.name, "chunk": i} for i in range(len(chunks))],
        )
        if tampilkan_log:
            print(f"  {pdf.name} → {len(chunks)} chunk (ukuran={ukuran}, overlap={tumpang})")


def indeks_dokumen(folder="dokumen"):
    """Lab 2 (Bagian 1) & Persiapan Lab 3 — Bangun indeks dari folder dokumen/."""
    folder_path = Path(folder)
    if not folder_path.exists():
        folder_path.mkdir(parents=True, exist_ok=True)
        print(f"Folder '{folder}/' telah dibuat. Silakan letakkan 2–3 file PDF di dalamnya.")
        return

    daftar_pdf = sorted(folder_path.glob("*.pdf"))
    if not daftar_pdf:
        print(f"Tidak ada file PDF di '{folder}/'. Silakan letakkan 2–3 file PDF di dalamnya.")
        return

    print("\n" + "=" * 70)
    print("LAB 2 — BANGUN INDEKS CHROMADB")
    print("=" * 70)
    print("[Koleksi Utama: 'pedoman' (ukuran=800, overlap=100)]")
    indeks_ke_koleksi(koleksi, daftar_pdf, ukuran=800, tumpang=100, tampilkan_log=True)

    print("\n[Koleksi Lab 3: 'pedoman_chunk_500' (ukuran=500, overlap=75)]")
    indeks_ke_koleksi(koleksi_kecil, daftar_pdf, ukuran=500, tumpang=75, tampilkan_log=True)

    print("\n[Koleksi Lab 3: 'pedoman_chunk_1500' (ukuran=1500, overlap=225)]")
    indeks_ke_koleksi(koleksi_besar, daftar_pdf, ukuran=1500, tumpang=225, tampilkan_log=True)


if __name__ == "__main__":
    eksperimen_chunking()
    indeks_dokumen()