import gradio as gr
from transformers import pipeline
from PIL import Image
import torch

# GPU болса қолдану, болмаса CPU.
device = "cuda:0" if torch.cuda.is_available() else "cpu"

# Модельдерді жүктеу үшін орын дайындау
waste_classifier = None
plant_classifier = None

# --- Қазақшаға аударуға арналған сөздіктер ---
# ⭐️ Қоқыс сөздігі (yangy50/garbage-classification-ға сәйкес қайтарылды)
WASTE_LABELS_KZ = {
    "cardboard": {"name": "📦 Картон", "recyclable": True, "description": "Бұл материал қайта өңдеуге жарамды. Таза және құрғақ күйінде арнайы қоқыс жәшігіне салыңыз."},
    "glass": {"name": "🍾 Шыны", "recyclable": True, "description": "Шыны бөтелкелер мен банкалар қайта өңделеді."},
    "metal": {"name": "🔩 Металл", "recyclable": True, "description": "Металл заттар қайта өңдеуге жатады."},
    "paper": {"name": "📄 Қағаз", "recyclable": True, "description": "Қағаздар қайта өңделеді. Майланған қағазды тастамаңыз."},
    "plastic": {"name": "🧴 Пластик", "recyclable": True, "description": "Пластик контейнерлер қайта өңделеді."},
    "trash": {"name": "🗑️ Басқа қоқыс", "recyclable": False, "description": "Бұл зат қайта өңдеуге жарамсыз."}
}

# ⭐️ Жаңа, Жеңіл Өсімдік моделіне сәйкес сөздік
PLANT_LABELS_KZ = {
    # Бұл модель PlantVillage-дің көптеген санаттарын қамтиды
    "Apple___Apple_scab": {"name": "🍎 Алманың қотыр ауруы", "description": "Жапырақтар мен жемістерде қара дақтар пайда болады."},
    "Apple___healthy": {"name": "🍏 Сау алма ағашы", "description": "Ауру белгілері жоқ. Өсімдік сау."},
    "Tomato___Bacterial_spot": {"name": "🍅 Қызанақтың бактериялық дағы", "description": "Сулы дақтар пайда болады."},
    "Tomato___Late_blight": {"name": "🍅 Қызанақтың кеш фитофторозы", "description": "Ірі қара дақтар пайда болады."},
    "Potato___Early_blight": {"name": "🥔 Картоптың ерте фитофторозы", "description": "Қара дақтар пайда болады."},
    "Potato___healthy": {"name": "🥔 Сау картоп", "description": "Өсімдікте ауру белгілері жоқ."},
    "Grape___healthy": {"name": "🍇 Сау жүзім", "description": "Өсімдік сау."},
    # Басқа аурулар модельдің өз жауабымен аударылады
}


# Суретті талдау функциясы
def analyze_image(mode, image):
    global waste_classifier, plant_classifier
    try:
        if mode == "♻️ Қоқыс":
            if waste_classifier is None:
                print("⭐️ Тұрақты Қоқыс тану моделін жүктеу: yangy50/garbage-classification")
                # ⭐️ СІЗ СҰРАҒАН ТҰРАҚТЫ МОДЕЛЬ
                waste_classifier = pipeline("image-classification", model="yangy50/garbage-classification", device=device)
            
            predictions = waste_classifier(image)
            top_prediction = predictions[0]
            label_en = top_prediction['label']
            
            result_kz = WASTE_LABELS_KZ.get(label_en, {"name": f"❓ {label_en.title()}", "recyclable": False, "description": "Анықтай алмадым."})
            status = "✅ Қайта өңделеді" if result_kz['recyclable'] else "❌ Өңделмейді"
            score_percent = top_prediction['score'] * 100
            
            return (f"**Нәтиже:** {result_kz['name']}\n\n"
                    f"**Статус:** {status}\n\n"
                    f"**Сипаттама:** {result_kz['description']}\n\n"
                    f"_(Сенімділік: {score_percent:.2f}%)_")

        elif mode == "🌿 Өсімдік":
            if plant_classifier is None:
                print("⭐️ Ең жеңіл және тұрақты Өсімдік моделін жүктеу: harman-thind/plant-disease-classification")
                # ⭐️ ЕҢ ЖЕҢІЛ БАЛАМА МОДЕЛЬ
                plant_classifier = pipeline("image-classification", model="harman-thind/plant-disease-classification", device=device)

            predictions = plant_classifier(image)
            top_prediction = predictions[0]
            # Модельдің жауап форматын тазалау
            label_en = top_prediction['label'].replace(" ", "_")
            
            result_kz = PLANT_LABELS_KZ.get(label_en, {"name": f"❓ {label_en.replace('_', ' ').title()}", "description": "Ауру түрін анықтай алмадым. Модельдің нәтижесі: " + label_en.replace('_', ' ')})
            score_percent = top_prediction['score'] * 100
            
            return (f"**Диагноз:** {result_kz['name']}\n\n"
                    f"**Сипаттама:** {result_kz['description']}\n\n"
                    f"_(Сенімділік: {score_percent:.2f}%)_")
            
    except Exception as e:
        # Егжей-тегжейлі қате туралы ақпарат
        return f"**Қате:** {type(e).__name__}: Модельді жүктеу кезінде мәселе туындады. Себебі: {str(e)}"

# --- Gradio интерфейсі (өзгеріссіз) ---
with gr.Blocks(theme=gr.themes.Soft(), title="AI Көмекші") as demo:
    gr.Markdown("## 🌌 AI Көмекші")
    
    with gr.Row():
        mode = gr.Radio(
            ["♻️ Қоқыс", "🌿 Өсімдік"],
            label="Режим",
            value="♻️ Қоқыс"
        )
    
    image_input = gr.Image(type="pil", label="Сурет")
    
    btn = gr.Button("🔍 Талдау", variant="primary")
        
    output_text = gr.Markdown(label="Нәтиже")
    
    btn.click(analyze_image, inputs=[mode, image_input], outputs=output_text)

demo.launch()