import streamlit as st
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch
import googletrans
from googletrans import Translator
import os
from groq import Groq

# Load Model
MODEL_NAME = "linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification"
processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)

# Groq API Key (Set in Hugging Face Secrets)
GROQ_API_KEY = os.getenv("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")
client = Groq(api_key=GROQ_API_KEY)

# Disease Descriptions
disease_info = {
    "Bacterial Spot": {"cause": "Bacteria (Xanthomonas spp.)", "remedy": "Use copper-based fungicides."},
    "Leaf Mold": {"cause": "Fungus (Cladosporium fulvum)", "remedy": "Use resistant plant varieties."},
    "Healthy": {"cause": "No disease detected", "remedy": "Your plant is healthy!"}
}

# Translator
translator = Translator()

# Streamlit UI
st.set_page_config(page_title="Plant Disease Detection", page_icon="🌿", layout="wide")

st.title("🌿 Plant Disease Detection App")
st.write("Upload a leaf image to detect diseases and get solutions.")

# Image Upload
uploaded_file = st.file_uploader("📷 Upload a leaf image...", type=["jpg", "png", "jpeg"])

if uploaded_file:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Predict Disease
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits

    predicted_class_idx = logits.argmax(-1).item()
    predicted_label = model.config.id2label[predicted_class_idx]
    confidence = torch.nn.functional.softmax(logits, dim=-1)[0][predicted_class_idx].item() * 100

    st.subheader(f"🔍 **Detected Disease:** {predicted_label} ({confidence:.2f}%)")

    # Get Disease Info
    if predicted_label in disease_info:
        cause = disease_info[predicted_label]["cause"]
        remedy = disease_info[predicted_label]["remedy"]
    else:
        cause = "Unknown cause."
        remedy = "Consult an expert."

    # Select Language
    language = st.selectbox("🌍 Select Language", list(googletrans.LANGUAGES.values()), index=21)  # Default: English

    lang_code = list(googletrans.LANGUAGES.keys())[list(googletrans.LANGUAGES.values()).index(language)]

    # Translate Disease Info
    cause_translated = translator.translate(cause, dest=lang_code).text
    remedy_translated = translator.translate(remedy, dest=lang_code).text

    st.info(f"🦠 **Cause:** {cause_translated}")
    st.success(f"💊 **Remedy:** {remedy_translated}")

    # Chatbot
    st.subheader("💬 Chat with AI about Plant Diseases")
    user_query = st.text_input("Type your question about the disease:")

    if user_query:
        response = client.chat.completions.create(
            messages=[{"role": "user", "content": user_query}],
            model="llama-3.3-70b-versatile"
        )
        chatbot_response = response.choices[0].message.content
        st.write("🤖 **Chatbot Response:**", chatbot_response)