import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch
import numpy as np
import logging
import requests
import os

# Configure Logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

# Load Model & Processor
model_name = "linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification"
try:
    processor = AutoImageProcessor.from_pretrained(model_name, use_fast=True)
    model = AutoModelForImageClassification.from_pretrained(model_name)
    logging.info("✅ Model and processor loaded successfully.")
except Exception as e:
    logging.error(f"❌ Failed to load model: {str(e)}")
    raise RuntimeError("Failed to load the model. Please check the logs for details.")

# Gemini API Key (Replace with your actual key)
GEMINI_API_KEY = os.getenv("AIzaSyCiRL0ES-zsJGJYsY03xmpwqcggDGcL2Fk", "AIzaSyCiRL0ES-zsJGJYsY03xmpwqcggDGcL2Fk")

# Function to Get AI-Powered Treatment Suggestions
def get_treatment_suggestions(disease_name):
    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateText?key={GEMINI_API_KEY}"
    headers = {"Content-Type": "application/json"}
    data = {
        "prompt": { "text": f"Provide detailed organic and chemical treatment options, including dosage and preventive care, for {disease_name} in crops." },
        "temperature": 0.7,
        "candidate_count": 1
    }

    try:
        response = requests.post(url, headers=headers, json=data)
        if response.status_code == 200:
            result = response.json()
            treatment = result.get("candidates", [{}])[0].get("content", "No treatment suggestions found.")
            return treatment
        else:
            logging.error(f"API Error: {response.status_code} - {response.text}")
            return f"API Error: {response.status_code}"
    except Exception as e:
        logging.error(f"Error fetching treatment suggestions: {str(e)}")
        return "Error retrieving treatment details."

# Define Prediction Function
def predict(image):
    try:
        image = Image.fromarray(np.uint8(image)).convert("RGB")
        inputs = processor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs)
        logits = outputs.logits
        predicted_class_idx = logits.argmax(-1).item()
        predicted_label = model.config.id2label[predicted_class_idx]
        
        # Get AI-generated treatment suggestions
        treatment = get_treatment_suggestions(predicted_label)
        
        return f"Predicted Disease: {predicted_label}\nTreatment: {treatment}"
    except Exception as e:
        logging.error(f"Prediction failed: {str(e)}")
        return f"❌ Prediction failed: {str(e)}"

# Gradio Interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Upload or capture plant image"),
    outputs=gr.Textbox(label="Result"),
    title="AI-Powered Plant Disease Detector",
    description="Upload a plant leaf image to detect diseases and get AI-powered treatment suggestions.",
    allow_flagging="never",
)

# Launch Gradio App
iface.launch()