import gradio as gr
from transformers import pipeline
import wikipedia

# Opsional: Mengatur pencarian Wikipedia ke bahasa Indonesia
wikipedia.set_lang("id")

# 1. LOAD MODEL KHUSUS TANAMAN
print("Sedang memuat model FloraSense...")
# Catatan: Jika model ini belum tersedia, Anda juga bisa memakai "google/vit-base-patch16-224"
classifier = pipeline("image-classification", model="Sisigoks/FloraSense")

def identify_plant(image):
    if image is None:
        return None, "⚠️ Silakan masukkan gambar terlebih dahulu."
    
    # Model melakukan prediksi
    predictions = classifier(image)
    
    # Format output untuk bar persentase (Progress Bar)
    confidences = {p["label"]: p["score"] for p in predictions}
    
    # Ambil label tanaman dengan probabilitas tertinggi untuk dicari di Wikipedia
    top_plant = predictions[0]["label"]
    
    # 2. PENCARIAN WIKIPEDIA
    try:
        # Mengambil 2 kalimat pertama dari Wikipedia
        wiki_summary = wikipedia.summary(top_plant, sentences=2)
        wiki_text = f"### 📖 Tahukah Anda?\n**{top_plant.title()}**: {wiki_summary}"
    except wikipedia.exceptions.DisambiguationError:
        wiki_text = f"### 📖 Info Wikipedia\nAda banyak arti untuk **{top_plant}**. Silakan cari secara spesifik."
    except wikipedia.exceptions.PageError:
        # Jika tidak ditemukan di Wikipedia Indonesia, beri pesan khusus
        wiki_text = f"### 📖 Info Wikipedia\nMaaf, informasi tentang **{top_plant}** belum tersedia di Wikipedia saat ini."
    except Exception:
        wiki_text = "### 📖 Info Wikipedia\nGagal terhubung ke server Wikipedia."

    return confidences, wiki_text

# 3. MEMBUAT ANTARMUKA GRADIO
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    <h1 style='text-align: center;'>🌿 Flora Lens: AI Plant Identifier</h1>
    <p style='text-align: center;'>Ambil foto tanaman dari kamera Anda, dan AI akan mendeteksi spesiesnya beserta info dari Wikipedia.</p>
    """)
    
    with gr.Row():
        with gr.Column():
            # PARAMETER BARU: sources=["webcam", "upload"] memunculkan icon kamera dan folder
            inp_img = gr.Image(type="pil", sources=["webcam", "upload"], label="📷 Ambil Foto / 📂 Upload Gambar")
            btn = gr.Button("🔍 Identifikasi & Cari Info", variant="primary")
            
        with gr.Column():
            out_label = gr.Label(num_top_classes=3, label="📊 Hasil Prediksi AI")
            out_wiki = gr.Markdown(label="Info Tambahan")
            
    # Menghubungkan tombol dengan fungsi (sekarang outputs ada dua: label dan teks wiki)
    btn.click(fn=identify_plant, inputs=inp_img, outputs=[out_label, out_wiki])

if __name__ == "__main__":
    demo.launch()