import chromadb
from chromadb.utils import embedding_functions


# ==========================================
# PENGATURAN
# ==========================================

FOLDER_DATABASE = "./chroma_db"

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


# ==========================================
# MEMBUAT KONEKSI
# ==========================================

def hubungkan_database(ukuran):
    client = chromadb.PersistentClient(
        path=FOLDER_DATABASE
    )

    embedding = (
        embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name=MODEL_EMBEDDING
        )
    )

    nama_koleksi = f"dokumen_chunk_{ukuran}"

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

    return koleksi


# ==========================================
# MELAKUKAN PENCARIAN
# ==========================================

def cari_dokumen():
    print("\n===== SISTEM PENCARIAN DOKUMEN RAG =====")
    print("Ukuran chunk tersedia: 500, 1000, 1500")

    try:
        ukuran = int(input("Pilih ukuran chunk: "))

        if ukuran not in [500, 1000, 1500]:
            print("Ukuran chunk tidak tersedia.")
            return

        koleksi = hubungkan_database(ukuran)

        jumlah_data = koleksi.count()

        if jumlah_data == 0:
            print("Indeks masih kosong. Jalankan indeks.py dahulu.")
            return

        print(f"\nJumlah chunk dalam indeks: {jumlah_data}")

        while True:
            pertanyaan = input(
                "\nMasukkan pertanyaan (ketik 'keluar' untuk berhenti): "
            ).strip()

            if pertanyaan.lower() == "keluar":
                break

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

            jumlah_hasil = min(3, jumlah_data)

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

            print("\n===== HASIL PENCARIAN =====")

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

            print(
                "Catatan: nilai jarak lebih kecil "
                "berarti hasil lebih dekat."
            )


    except ValueError:
        print("Masukkan angka ukuran chunk yang benar.")

    except Exception as error:
        print("Terjadi kesalahan:", error)
        print("Pastikan indeks sudah dibuat dengan indeks.py.")


# ==========================================
# PROGRAM UTAMA
# ==========================================

if __name__ == "__main__":
    cari_dokumen()