import os
# Force the use of legacy Keras (Keras 2 behavior) so that hub.KerasLayer is recognized properly.
os.environ["TF_USE_LEGACY_KERAS"] = "1"

import streamlit as st
import tensorflow as tf
import tensorflow_hub as hub
import numpy as np
from PIL import Image
import requests
from googletrans import Translator
import asyncio
import nest_asyncio

# Allow nested event loops in Streamlit
nest_asyncio.apply()

st.set_page_config(
    page_title="Plant Disease Classifier 🌱",
    page_icon="🌱",
    layout="wide",
    initial_sidebar_state="expanded",
)

# Custom CSS for styling
custom_css = """
<style>
body { background-color: #f8f9fa; }
h1, h2, h3, h4 { color: #2c3e50; font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; }
.stButton>button { background-color: #27ae60; color: white; border: none; padding: 0.5em 1em; border-radius: 5px; font-size: 16px; }
.sidebar .sidebar-content { background-image: linear-gradient(#27ae60, #2ecc71); color: white; }
</style>
"""
st.markdown(custom_css, unsafe_allow_html=True)

# Dictionary mapping diseases to recommended pesticides
pesticide_recommendations = {
    'Bacterial Blight': 'Copper-based fungicides, Streptomycin',
    'Red Rot': 'Fungicides containing Mancozeb or Copper',
    'Blight': 'Fungicides containing Chlorothalonil',
    'Common_Rust': 'Fungicides containing Azoxystrobin or Propiconazole',
    'Gray_Leaf_Spot,Healthy': 'Fungicides containing Azoxystrobin or Propiconazole',
    'Bacterial blight': 'Copper-based fungicides, Streptomycin',
    'curl_virus': 'Insecticides such as Imidacloprid or Pyrethroids',
    'fussarium_wilt': 'Soil fumigants, Fungicides containing Thiophanate-methyl',
    'Bacterial_blight': 'Copper-based fungicides, Streptomycin',
    'Blast': 'Fungicides containing Tricyclazole or Propiconazole',
    'Brownspot': 'Fungicides containing Azoxystrobin or Propiconazole',
    'Tungro': 'Insecticides such as Neonicotinoids or Pyrethroids',
    'septoria': 'Fungicides containing Azoxystrobin or Propiconazole',
    'strip_rust': 'Fungicides containing Azoxystrobin or Propiconazole'
}

def recommend_pesticide(predicted_class):
    if predicted_class == 'Healthy':
        return 'No need for any pesticide, plant is healthy'
    return pesticide_recommendations.get(predicted_class, "No recommendation available")

# Use st.cache_resource to load H5 models.
@st.cache_resource(show_spinner=False)
def load_h5_model(model_path):
    # This will load your H5 Keras model (which contains hub.KerasLayer)
    return tf.keras.models.load_model(
        model_path,
        custom_objects={"KerasLayer": hub.KerasLayer},
        compile=False
    )

# Define your models dictionary (update paths as needed)
models = {
    'sugarcane': load_h5_model("models/sugercane_model.h5"),
    'maize': load_h5_model("models/maize_model.h5"),
    'cotton': load_h5_model("models/cotton_model.h5"),
    'rice': load_h5_model("models/rice.h5"),
    'wheat': load_h5_model("models/wheat_model.h5"),
}

# Class names for each model (ensure these match the order of your model outputs)
class_names = {
    'sugarcane': ['Bacterial Blight', 'Healthy', 'Red Rot'],
    'maize': ['Blight', 'Common_Rust', 'Gray_Leaf_Spot,Healthy'],
    'cotton': ['Bacterial blight', 'curl_virus', 'fussarium_wilt', 'Healthy'],
    'rice': ['Bacterial_blight', 'Blast', 'Brownspot', 'Tungro'],
    'wheat': ['Healthy', 'septoria', 'strip_rust'],
}

# Preprocess the uploaded image
def preprocess_image(image_file):
    try:
        image = Image.open(image_file).convert("RGB")
        image = image.resize((224, 224))
        img_array = np.array(image).astype("float32") / 255.0
        return np.expand_dims(img_array, axis=0)
    except Exception as e:
        st.error("Error processing image. Please upload a valid image file.")
        return None

# Classify the image using the appropriate model
def classify_image(model_name, image_file):
    input_image = preprocess_image(image_file)
    if input_image is None:
        return None, None
    predictions = models[model_name].predict(input_image)
    predicted_index = np.argmax(predictions)
    predicted_class = class_names[model_name][predicted_index]
    recommended_pesticide = recommend_pesticide(predicted_class)
    return predicted_class, recommended_pesticide

# (Optional) Retrieve detailed plant information from LM Studio
def get_plant_info(disease, plant_type="Unknown"):
    prompt = f"""
    Disease Name: {disease}
    Plant Type: {plant_type}

    Explain this disease in a very simple way for a farmer. Include:
    - Symptoms
    - Causes
    - Severity
    - How It Spreads
    - Treatment & Prevention
    """
    try:
        LM_STUDIO_API_URL = os.getenv("LM_STUDIO_API_URL", "http://192.168.56.1:1234/v1/chat/completions")
        response = requests.post(LM_STUDIO_API_URL, json={"messages": [{"role": "user", "content": prompt}]})
        response.raise_for_status()
        data = response.json()
        detailed_info = data.get("choices", [{}])[0].get("message", {}).get("content", "")
        return {"detailed_info": detailed_info}
    except Exception as e:
        st.error("Error retrieving detailed plant info.")
        return {"detailed_info": ""}

def get_web_pesticide_info(disease, plant_type="Unknown"):
    query = f"site:agrowon.esakal.com {disease} in {plant_type}"
    url = "https://www.googleapis.com/customsearch/v1"
    params = {
        "key": os.getenv("GOOGLE_API_KEY"),
        "cx": os.getenv("GOOGLE_CX"),
        "q": query,
        "num": 3
    }
    try:
        response = requests.get(url, params=params)
        response.raise_for_status()
        data = response.json()
        if "items" in data and len(data["items"]) > 0:
            item = data["items"][0]
            return {"title": item.get("title", "No title available"),
                    "link": item.get("link", "#"),
                    "snippet": item.get("snippet", "No snippet available"),
                    "summary": item.get("snippet", "No snippet available")}
    except Exception as e:
        st.error("Error retrieving web pesticide info.")
    return None

def get_more_web_info(query):
    url = "https://www.googleapis.com/customsearch/v1"
    params = {
        "key": os.getenv("GOOGLE_API_KEY"),
        "cx": os.getenv("GOOGLE_CX"),
        "q": query,
        "num": 3
    }
    try:
        response = requests.get(url, params=params)
        response.raise_for_status()
        data = response.json()
        results = []
        if "items" in data:
            for item in data["items"]:
                results.append({"title": item.get("title", "No title available"),
                                "link": item.get("link", "#"),
                                "snippet": item.get("snippet", "No snippet available")})
        return results
    except Exception as e:
        st.error("Error retrieving additional articles.")
    return []

def get_commercial_product_info(recommendation):
    indiamart_query = f"site:indiamart.com pesticide '{recommendation}'"
    krishi_query = f"site:krishisevakendra.in/products pesticide '{recommendation}'"
    indiamart_results = get_more_web_info(indiamart_query)
    krishi_results = get_more_web_info(krishi_query)
    return indiamart_results + krishi_results

translator = Translator()
if "language" not in st.session_state:
    st.session_state.language = "English"

async def async_translate_text(text):
    translated = await translator.translate(text, src='en', dest='mr')
    return translated.text

def translate_text(text):
    if st.session_state.get("language", "English") == "Marathi":
        try:
            return asyncio.get_event_loop().run_until_complete(async_translate_text(text))
        except Exception as e:
            st.error("Translation error.")
            return text
    return text

def main():
    st.sidebar.title("Settings")
    st.sidebar.info("Choose language and plant type, then upload an image to classify the disease.")
    language_option = st.sidebar.radio("Language", options=["English", "Marathi"], index=0)
    st.session_state.language = language_option
    plant_type = st.sidebar.selectbox("Select Plant Type", options=['sugarcane', 'maize', 'cotton', 'rice', 'wheat'])
    uploaded_file = st.sidebar.file_uploader("Upload a plant image...", type=["jpg", "jpeg", "png"])

    col1, col2 = st.columns([1, 2])
    with col1:
        st.image("https://via.placeholder.com/150.png?text=Plant", caption=translate_text("Plant Health"), use_container_width=True)
    with col2:
        st.title(translate_text("Krushi Mitra"))
        st.write(translate_text("Plant Disease Classification and Pesticide Recommendation.\n\nUpload an image, select plant type, and click on Classify."))

    if uploaded_file is not None:
        st.markdown("---")
        st.subheader(translate_text("Uploaded Image"))
        st.image(uploaded_file, use_container_width=True)
        if st.button(translate_text("Classify")):
            with st.spinner(translate_text("Classifying...")):
                predicted_class, pesticide = classify_image(plant_type, uploaded_file)
            if predicted_class:
                st.success(translate_text("Classification Complete!"))
                st.markdown(f"### {translate_text('Predicted Class')} ({plant_type.capitalize()}): {translate_text(predicted_class)}")
                st.markdown(f"### {translate_text('Recommended Pesticide')}: {translate_text(pesticide)}")
                tabs = st.tabs([translate_text("Detailed Info"), translate_text("Commercial Products"), translate_text("More Articles")])
                with tabs[0]:
                    with st.spinner(translate_text("Retrieving detailed plant information...")):
                        info = get_plant_info(predicted_class, plant_type)
                    if info and info.get("detailed_info"):
                        st.markdown(translate_text("#### Detailed Plant Disease Information"))
                        st.markdown(translate_text(info.get("detailed_info")))
                    else:
                        st.info(translate_text("Detailed information is not available at the moment."))
                    web_recommendation = get_web_pesticide_info(predicted_class, plant_type)
                    if web_recommendation:
                        st.markdown(translate_text("#### Additional Pesticide Recommendations"))
                        st.markdown(f"**{translate_text('Title')}:** {translate_text(web_recommendation['title'])}")
                        st.markdown(f"**{translate_text('Summary')}:** {translate_text(web_recommendation['summary'])}")
                        if web_recommendation['link']:
                            st.markdown(f"[{translate_text('Read More')}]({web_recommendation['link']})")
                    else:
                        st.info(translate_text("No additional pesticide recommendations available."))
                with tabs[1]:
                    with st.spinner(translate_text("Retrieving commercial product details...")):
                        commercial_products = get_commercial_product_info(pesticide)
                    if commercial_products:
                        for item in commercial_products:
                            st.markdown(f"**{translate_text('Title')}:** {translate_text(item['title'])}")
                            st.markdown(f"**{translate_text('Snippet')}:** {translate_text(item['snippet'])}")
                            if item['link']:
                                st.markdown(f"[{translate_text('Read More')}]({item['link']})")
                            st.markdown("---")
                    else:
                        st.info(translate_text("No commercial product details available."))
                with tabs[2]:
                    with st.spinner(translate_text("Retrieving additional articles...")):
                        more_info = get_more_web_info(f"{predicted_class} in {plant_type}")
                    if more_info:
                        for item in more_info:
                            st.markdown(f"**{translate_text('Title')}:** {translate_text(item['title'])}")
                            st.markdown(f"**{translate_text('Snippet')}:** {translate_text(item['snippet'])}")
                            if item['link']:
                                st.markdown(f"[{translate_text('Read More')}]({item['link']})")
                            st.markdown("---")
                    else:
                        st.info(translate_text("No additional articles available."))
            else:
                st.error(translate_text("Error in classification. Please try again."))
    else:
        st.info(translate_text("Please upload an image from the sidebar to get started."))

if __name__ == "__main__":
    main()