import os
import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image

import google.generativeai as genai
from dotenv import load_dotenv

# ✅ Load Gemini API Key from environment
load_dotenv()
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))

# ✅ Load trained model
model = tf.keras.models.load_model("model.h5")
IMG_SIZE = (64, 64)

# ✅ Class labels (10 diseases)
class_names = [
    "bacterial_leaf_blight",
    "bacterial_leaf_streak",
    "bacterial_panicle_blight",
    "blast",
    "brown_spot",
    "dead_heart",
    "downy_mildew",
    "hispa",
    "normal",
    "tungro"
]

# ✅ Get Gemini guidance in English (clean formatting)
def get_gemini_diagnosis(disease_label: str) -> str:
    prompt = f"""
You are a plant pathology expert and agricultural advisor.

A rice plant is infected with: {disease_label}.
Give a short, farmer-friendly explanation with:
1. Causes and symptoms
2. Treatments or preventive actions
3. Risk level for crops

Avoid medical jargon. Keep it clear, helpful, and practical.
"""
    model = genai.GenerativeModel("gemini-2.5-flash")
    response = model.generate_content(prompt)
    return response.text.strip().replace("*", "")  # ✅ Remove asterisks for clean UI


# ✅ Translate Gemini response into selected language
def translate_to_language(text: str, language: str) -> str:
    if language.lower() == "english":
        return text  # No need to translate

    prompt = f"Translate the following expert guidance into {language}:\n\n{text}"
    model = genai.GenerativeModel("gemini-2.5-flash")
    response = model.generate_content(prompt)
    return response.text.strip()


# ✅ Main pipeline: prediction + Gemini + translation
def full_diagnosis(image, language):
    image = image.convert("RGB").resize(IMG_SIZE)
    img_array = np.array(image) / 255.0
    img_array = np.expand_dims(img_array, axis=0)

    predictions = model.predict(img_array)[0]
    top_idx = np.argmax(predictions)
    top_label = class_names[top_idx]
    confidence = float(predictions[top_idx]) * 100

    # Get English response first
    gemini_response = get_gemini_diagnosis(top_label)

    # Translate only if Urdu or Hindi selected
    final_response = translate_to_language(gemini_response, language)

    return f"{top_label} ({confidence:.2f}% confidence)", final_response


# ✅ Gradio UI with Language Dropdown and Dark/Light Theme Toggle
with gr.Blocks(theme=gr.themes.Base()) as demo:
    gr.Markdown("## 🌾 CropDoctor – Rice Disease Detector with Gemini AI")
    gr.Markdown("Upload a rice leaf image → Model predicts disease → Gemini AI gives expert guidance")

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Rice Leaf Image")
        lang_dropdown = gr.Dropdown(
            choices=["English", "Urdu", "Hindi"],
            value="English",
            label="Choose Language for AI Guidance"
        )

    output1 = gr.Textbox(label="Predicted Disease")
    output2 = gr.Textbox(label="AI Expert Guidance (Translated)")

    btn = gr.Button("Diagnose")
    btn.click(fn=full_diagnosis, inputs=[image_input, lang_dropdown], outputs=[output1, output2])

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