import chromadb
from chromadb.utils import embedding_functions

# Menggunakan model embedding yang sama
emb = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="paraphrase-multilingual-MiniLM-L12-v2"
)

# Membuka database
client = chromadb.PersistentClient(path="./chroma_db")

# Membuka koleksi ukuran 1000
koleksi = client.get_collection(
    name="pedoman_1000",
    embedding_function=emb
)

# 5 pertanyaan biasa + 2 parafrase
PERTANYAAN = [
    ("Biasa", "Apa yang dimaksud dengan airworthiness?"),
    ("Biasa", "Apa kepanjangan dan pengertian MRO?"),
    ("Biasa", "Apa perbedaan line maintenance dan base maintenance?"),
    ("Biasa", "Apa tujuan utama aircraft maintenance?"),
    ("Biasa", "Apa itu Non-Destructive Testing?"),

    ("Parafrase", "Bagaimana mengetahui bahwa sebuah pesawat layak diterbangkan dengan aman?"),
    ("Parafrase", "Jenis perawatan rutin apa yang dilakukan ketika pesawat masih beroperasi?")
]

if koleksi.count() == 0:
    print("Database masih kosong.")

else:
    for nomor, (jenis, pertanyaan) in enumerate(PERTANYAAN, 1):

        print("\n" + "=" * 60)
        print(f"Pertanyaan {nomor} ({jenis})")
        print(pertanyaan)

        hasil = koleksi.query(
            query_texts=[pertanyaan],
            n_results=min(3, koleksi.count())
        )

        for peringkat, (doc, meta, jarak) in enumerate(
            zip(
                hasil["documents"][0],
                hasil["metadatas"][0],
                hasil["distances"][0]
            ),
            start=1
        ):
            print(f"\nPeringkat {peringkat}")
            print(f"Sumber: {meta['sumber']}")
            print(f"Chunk: {meta['chunk']}")
            print(f"Jarak: {jarak:.3f}")
            print(doc[:300])