import gradio as gr
from ultralytics import YOLO
from collections import Counter
import os
import uuid

# Model yükle
model = YOLO("best.pt")

# Sınıf eşlemesi
id2name = {
    0: "Silah", 1: "Tank", 2: "Helikopter", 3: "Savaş Uçağı", 4: "Asker",
    5: "Top", 6: "Füze", 7: "Savaş Gemisi", 8: "Askeri Araç", 9: "Bayrak",
    10: "Sivil Araç", 11: "Ağaç"
}

PRED_DIR = "yolo_outputs"

def save_and_predict(image):
    os.makedirs("inputs", exist_ok=True)
    image_id = str(uuid.uuid4())[:8]
    image_path = f"inputs/{image_id}.jpg"
    image.save(image_path)

    model.predict(source=image_path, save_txt=True, conf=0.25, project=PRED_DIR, name=image_id, exist_ok=True)

    txt_path = f"{PRED_DIR}/{image_id}/labels/{image_id}.txt"

    if not os.path.exists(txt_path):
        return {}, "Hiçbir nesne tespit edilemedi."

    with open(txt_path, "r") as f:
        class_ids = [int(line.strip().split()[0]) for line in f]

    counts = Counter(class_ids)
    detected_names = {id2name[i] for i in counts if i in id2name}

    data = {
        "counts": counts,
        "names": detected_names
    }

    return data, "Tahmin başarıyla yüklendi. Sağdaki kutulardan bir soruya tıklayın."

def cevapla(soru_id, data, chat_history):
    if not data:
        return chat_history + [("❓ " + soru_id, "Önce bir görsel yükleyin ve tahmin yapın.")]

    counts = data["counts"]
    names = data["names"]

    if soru_id == "1️⃣ Kaç nesne ve kaçar adet?":
        detaylar = [f"{counts[i]} adet {id2name[i]}" for i in counts]
        cevap = f"{len(counts)} farklı sınıf tespit edildi:\n" + "\n".join(detaylar)

    elif soru_id == "2️⃣ Savunma sistemleri var mı?":
        savunma = {"Tank", "Füze", "Top", "Helikopter"}
        bulunan = savunma & names
        cevap = f"Bulunan savunma sistemleri: {', '.join(bulunan)}" if bulunan else "Savunma sistemi tespit edilmedi."

    elif soru_id == "3️⃣ Askeri birlik tespiti?":
        cevap = "Askeri birlik olabilir." if {"Asker", "Askeri Araç"} & names else "Askeri birlik tespit edilmedi."

    elif soru_id == "4️⃣ Hava unsurları var mı?":
        hava = {"Helikopter", "Savaş Uçağı"}
        bulunan = hava & names
        cevap = f"Hava unsurları tespit edildi: {', '.join(bulunan)}" if bulunan else "Hava unsuru tespit edilmedi."

    elif soru_id == "5️⃣ Tehdit seviyesi nedir?":
        yüksek = {"Tank", "Füze", "Savaş Uçağı", "Top", "Savaş Gemisi"}
        if yüksek & names:
            cevap = "Yüksek tehdit içeriyor."
        elif "Asker" in names:
            cevap = "Orta seviye tehdit."
        else:
            cevap = "Düşük tehdit seviyesi."

    else:
        cevap = "Geçersiz soru."

    return chat_history + [(soru_id, cevap)]

# Gradio Blocks yapısı
with gr.Blocks() as demo:
    gr.Markdown("## 🪖 YOLO Görselden Nesne Tespiti + Soru-Cevap Arayüzü")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Görsel Yükle")
            tahmin_bilgisi = gr.Textbox(label="Tahmin Durumu")
            hidden_data = gr.State()
            upload_button = gr.Button("📤 Tahmin Yap")
        with gr.Column():
            chatbot = gr.Chatbot(label="📡 Soru-Cevap Paneli")
            with gr.Row():
                b1 = gr.Button("1️⃣ Kaç nesne ve kaçar adet?")
                b2 = gr.Button("2️⃣ Savunma sistemleri var mı?")
            with gr.Row():
                b3 = gr.Button("3️⃣ Askeri birlik tespiti?")
                b4 = gr.Button("4️⃣ Hava unsurları var mı?")
            b5 = gr.Button("5️⃣ Tehdit seviyesi nedir?")

    # Buton olayları
        # Buton olayları – dikkat: soru_id doğrudan lambda ile veriliyor
    upload_button.click(
        save_and_predict, inputs=image_input, outputs=[hidden_data, tahmin_bilgisi]
    )

    b1.click(lambda data, chat: cevapla("1️⃣ Kaç nesne ve kaçar adet?", data, chat), inputs=[hidden_data, chatbot], outputs=chatbot)
    b2.click(lambda data, chat: cevapla("2️⃣ Savunma sistemleri var mı?", data, chat), inputs=[hidden_data, chatbot], outputs=chatbot)
    b3.click(lambda data, chat: cevapla("3️⃣ Askeri birlik tespiti?", data, chat), inputs=[hidden_data, chatbot], outputs=chatbot)
    b4.click(lambda data, chat: cevapla("4️⃣ Hava unsurları var mı?", data, chat), inputs=[hidden_data, chatbot], outputs=chatbot)
    b5.click(lambda data, chat: cevapla("5️⃣ Tehdit seviyesi nedir?", data, chat), inputs=[hidden_data, chatbot], outputs=chatbot)


demo.launch()