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="Krushi Mitra 🌱",
        page_icon="🌱",
        layout="wide",  # Change to centered layout
        initial_sidebar_state="auto"  # Auto-hide sidebar on small screens
)

# Improved translations dictionary
TRANSLATIONS = {
    # Common Website Words
    "Upload": "अपलोड करा",
    "Select": "निवडा",
    "Image": "प्रतिमा",
    "Settings": "सेटिंग्ज",
    "Language": "भाषा",
    "Plant": "पीक",
    "Disease": "रोग",
    "Classify": "वर्गीकृत करा",
    "Recommended": "शिफारस केलेला",
    "More": "अधिक",
    "Read": "वाचा",
    
    # Specific Phrases
    "Krushi Mitra": "कृषी मित्र",
    "Plant Disease Classification and Pesticide Recommendation": "पीक रोग वर्गीकरण आणि कीटकनाशक शिफारस",
    "Upload a plant image...": "पीक प्रतिमा अपलोड करा...",
    "Select an image of your plant to detect diseases": "रोग शोधण्यासाठी आपल्या पीकाची प्रतिमा निवडा",
    "Uploaded Image": "अपलोड केलेली प्रतिमा",
    "Classifying...": "वर्गीकरण करत आहे...",
    "Classification Complete!": "वर्गीकरण पूर्ण झाले!",
    "Predicted Class": "अनुमानित वर्ग",
    "Recommended Pesticide": "शिफारस केलेला कीटकनाशक",
    "Detailed Info": "विस्तृत माहिती",
    "Commercial Products": "व्यावसायिक उत्पादने",
    "More Articles": "अधिक लेख",
    "Retrieving detailed plant information...": "पीकाची विस्तृत माहिती मिळवत आहे...",
    "Detailed Plant Disease Information": "पीक रोगाची विस्तृत माहिती",
    "Detailed information is not available at the moment.": "सध्या विस्तृत माहिती उपलब्ध नाही.",
    "Additional Pesticide Recommendations": "अतिरिक्त कीटकनाशक शिफारसी",
    "Title": "शीर्षक",
    "Summary": "सारांश",
    "Retrieving commercial product details...": "व्यावसायिक उत्पादन तपशील मिळवत आहे...",
    "No commercial product details available.": "कोणतेही व्यावसायिक उत्पादन तपशील उपलब्ध नाहीत.",
    "Retrieving additional articles...": "अतिरिक्त लेख मिळवत आहे...",
    "No additional articles available.": "अतिरिक्त लेख उपलब्ध नाहीत.",
    "Error in classification. Please try again.": "वर्गीकरणमध्ये त्रुटी. कृपया पुन्हा प्रयत्न करा.",
    "Please upload an image from the sidebar to get started.": "सुरुवात करण्यासाठी कृपया बाजूच्या पट्टीतून एक प्रतिमा अपलोड करा."
}

def translate_text(text):
    if st.session_state.get("language", "English") == "Marathi":
        return TRANSLATIONS.get(text, text)
    return text

# Custom CSS for styling
# Update custom CSS
custom_css = """
<style>
body { 
    background-color: #f0f2f6; 
    font-family: 'Inter', sans-serif; 
}
.stApp {  
    margin: 0 auto; 
}
.stButton>button { 
    background-color: #2ecc71; 
    color: white; 
    border: none; 
    padding: 10px 20px; 
    border-radius: 8px; 
    font-weight: bold; 
    transition: all 0.3s ease;
}
.stButton>button:hover {
    background-color: #27ae60;
    transform: scale(1.05);
}
.stTabs [data-baseweb="tab-list"] {
    background-color: #ecf0f1;
    border-radius: 10px;
}
.stTabs [data-baseweb="tab"] {
    padding: 10px 15px;
    font-weight: 600;
}
</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:
        API_URL = "https://api-inference.huggingface.co/models/meta-llama/Llama-3.2-1B-Instruct/v1/chat/completions"
        headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACE_API_TOKEN')}"}
        try:
            data = {
            "messages": [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt},
            ]
        }
            response = requests.post(API_URL, headers=headers, json=data)
            response.raise_for_status()
            print("Response:", data)
        except requests.exceptions.HTTPError as http_err:
            print("HTTP error occurred:", http_err)
            print("Response Content:", response.content)
        except Exception as e:
            print("Other error occurred:", str(e))
        response_data  = response.json()
        detailed_info = response_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(translate_text("Settings"))
    st.sidebar.info(translate_text("Choose language and plant type, then upload an image to classify the disease."))
    language_option = st.sidebar.radio(translate_text("Language"), options=["English", "Marathi"], index=0)
    st.session_state.language = language_option
    plant_type = st.sidebar.selectbox(translate_text("Select Plant Type"), options=['sugarcane', 'maize', 'cotton', 'rice', 'wheat'])
    uploaded_file = st.sidebar.file_uploader(
        translate_text("Upload a plant image..."), 
        type=["jpg", "jpeg", "png"],
        help=translate_text("Select an image of your plant to detect diseases")
    )
    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, width=300)
        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()