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

# 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.")

# Function to Fetch Treatment Suggestions from the Internet
def fetch_treatment_info(disease_name):
    try:
        search_url = f"https://www.bing.com/search?q=treatment+for+{disease_name.replace(' ', '+')}+in+plants"
        headers = {"User-Agent": "Mozilla/5.0"}
        response = requests.get(search_url, headers=headers)

        if response.status_code == 200:
            soup = BeautifulSoup(response.text, "html.parser")
            snippets = soup.find_all("p")  
            treatments = [s.text for s in snippets][:3]  

            return "\n".join(treatments) if treatments else "No treatment suggestions found."
        else:
            return "Failed to fetch treatment suggestions."
    except Exception as e:
        logging.error(f"Error fetching treatment info: {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()
        confidence = torch.nn.functional.softmax(logits, dim=-1)[0][predicted_class_idx] * 100
        predicted_label = model.config.id2label[predicted_class_idx]

        treatment = fetch_treatment_info(predicted_label)

        return predicted_label, confidence.item(), treatment
    except Exception as e:
        logging.error(f"Prediction failed: {str(e)}")
        return "Error", 0, "Error in prediction."

# Create a Gradio UI with Tabs & Styling
with gr.Blocks(css="body {background-color: #f7f9fc;}") as app:
    gr.Markdown(
        "<h1 style='text-align: center;'>🌿 AI-Powered Plant Disease Detector</h1>",
        elem_id="title"
    )
    
    with gr.Tabs():
        with gr.TabItem("📷 Detect Disease"):
            with gr.Row():
                image_input = gr.Image(type="numpy", label="📸 Upload a plant image", interactive=True)
            
            with gr.Row():
                btn_predict = gr.Button("🔍 Analyze", variant="primary")
            
            with gr.Row():
                disease_output = gr.Textbox(label="🌱 Disease", interactive=False)
                confidence_output = gr.Slider(minimum=0, maximum=100, interactive=False, label="🔬 Confidence (%)")
            
            with gr.Row():
                treatment_output = gr.Textbox(label="💊 Suggested Treatment", interactive=False)
            
            btn_predict.click(predict, inputs=[image_input], outputs=[disease_output, confidence_output, treatment_output])
        
        with gr.TabItem("📖 Prevention Guide"):
            gr.Markdown("""
                ## 🌾 How to Keep Your Plants Healthy?
                - 🏡 Ensure Proper Spacing to Avoid Fungal Growth
                - 💧 Water in the Morning to Reduce Disease Spread
                - 🛑 Regularly Inspect & Remove Infected Leaves
                - ☀️ Allow Good Sunlight & Ventilation for Stronger Plants
                - 🍀 Use Organic Pesticides & Fungicides When Necessary
            """)
    
app.launch()