import os
import gradio as gr
import numpy as np
import requests
import json
from dotenv import load_dotenv
from tensorflow.keras.models import load_model
from PIL import Image

# Load environment variables
load_dotenv()

# ===== Groq API Key =====
GROQ_API_KEY = os.getenv("GROQ_API_KEY", "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")
GROQ_MODEL = "llama3-70b-8192"  # Using Llama 3 70B model
print(f"Groq API key available: {'Yes' if GROQ_API_KEY else 'No'}")

# ===== Fallback to Hugging Face API Token =====
HF_API_TOKEN = os.getenv("HF_API_TOKEN")
print(f"HF API token available: {'Yes' if HF_API_TOKEN else 'No'}")

# ===== Load Trained Models =====
model_a = load_model("Tomato_Leaf_Disease_Model.h5")
model_b = load_model("tomato_leaf_model_final(77%).h5")
classifier_model = load_model("tomato_leaf_classifier_optimized.h5")

# ===== Disease Information Database (fallback if API fails) =====
disease_info = {
    "Tomato Bacterial Spot": {
        "description": "A bacterial disease that causes small, dark spots on leaves, stems, and fruits.",
        "causes": "Caused by Xanthomonas bacteria, spread by water splash, contaminated tools, and seeds.",
        "treatment": [
            "Remove and destroy infected plants",
            "Rotate crops with non-solanaceous plants",
            "Use copper-based fungicides",
            "Avoid overhead irrigation"
        ]
    },
    "Tomato Early Blight": {
        "description": "A fungal disease that causes dark spots with concentric rings on lower leaves first.",
        "causes": "Caused by Alternaria solani fungus, favored by warm, humid conditions.",
        "treatment": [
            "Remove infected leaves promptly",
            "Improve air circulation around plants",
            "Apply fungicides preventatively",
            "Mulch around plants to prevent soil splash"
        ]
    },
    "Tomato Late Blight": {
        "description": "A devastating fungal disease that causes dark, water-soaked lesions on leaves and fruits.",
        "causes": "Caused by Phytophthora infestans, favored by cool, wet conditions.",
        "treatment": [
            "Remove and destroy infected plants immediately",
            "Apply fungicides preventatively in humid conditions",
            "Improve drainage and air circulation",
            "Plant resistant varieties when available"
        ]
    },
    "Tomato Mosaic Virus": {
        "description": "A viral disease that causes mottled green/yellow patterns on leaves and stunted growth.",
        "causes": "Caused by tobacco mosaic virus (TMV), spread by handling, tools, and sometimes seeds.",
        "treatment": [
            "Remove and destroy infected plants",
            "Wash hands and tools after handling infected plants",
            "Control insect vectors like aphids",
            "Plant resistant varieties"
        ]
    },
    "Tomato Yellow Leaf Curl Virus": {
        "description": "A viral disease transmitted by whiteflies that causes yellowing and curling of leaves.",
        "causes": "Caused by a begomovirus, transmitted primarily by whiteflies.",
        "treatment": [
            "Use whitefly control measures",
            "Remove and destroy infected plants",
            "Use reflective mulches to repel whiteflies",
            "Plant resistant varieties"
        ]
    },
    "Tomato___Target_Spot": {
        "description": "A fungal disease causing circular lesions with concentric rings on leaves, stems, and fruits.",
        "causes": "Caused by Corynespora cassiicola fungus, favored by warm, humid conditions.",
        "treatment": [
            "Remove infected plant parts",
            "Improve air circulation",
            "Apply fungicides at first sign of disease",
            "Avoid overhead irrigation"
        ]
    },
    "Tomato___Bacterial_spot": {
        "description": "A bacterial disease causing small, dark, water-soaked spots on leaves, stems, and fruits.",
        "causes": "Caused by Xanthomonas species, spread by water splash and contaminated tools.",
        "treatment": [
            "Remove infected plant debris",
            "Use copper-based bactericides",
            "Rotate crops",
            "Use disease-free seeds and transplants"
        ]
    },
    "Tomato___healthy": {
        "description": "The plant shows no signs of disease and appears to be in good health.",
        "causes": "Proper growing conditions, good management practices, and disease prevention.",
        "treatment": [
            "Continue regular watering and fertilization",
            "Monitor for early signs of disease",
            "Maintain good air circulation",
            "Practice crop rotation"
        ]
    }
}

# ===== Preprocessing Function =====
def preprocess_image(image, target_size=(224, 224)):
    # Ensure the image is resized and normalized.
    if isinstance(image, Image.Image):
        img = image.resize(target_size)
    else:
        img = Image.fromarray(image).resize(target_size)
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

# ===== Disease Label Mappings =====
# Model A labels
disease_labels_a = {
    0: "Tomato Bacterial Spot",
    1: "Tomato Early Blight",
    2: "Tomato Late Blight",
    3: "Tomato Mosaic Virus",
    4: "Tomato Yellow Leaf Curl Virus"
}

# Model B labels
disease_labels_b = {
    0: "Tomato___Target_Spot",
    1: "Tomato___Bacterial_spot",
    2: "Tomato___Early_blight",
    3: "Tomato___healthy",
    4: "Tomato___Late_blight"
}

# ===== Prediction Functions =====
def predict_model_a(image):
    img = preprocess_image(image)
    pred = model_a.predict(img)
    predicted_class = np.argmax(pred)
    confidence = float(np.max(pred) * 100)
    return disease_labels_a.get(predicted_class, "Unknown result"), confidence

def predict_model_b(image):
    img = preprocess_image(image)
    pred = model_b.predict(img)
    predicted_class = np.argmax(pred)
    confidence = float(np.max(pred) * 100)
    return disease_labels_b.get(predicted_class, "Unknown result"), confidence

def predict_classifier(image):
    img = preprocess_image(image)
    pred = classifier_model.predict(img)
    # Here we assume the classifier returns class 1 for "Tomato Leaf"
    return "Tomato Leaf" if np.argmax(pred) == 1 else "Not Tomato Leaf"

# ===== Groq API Call =====
def call_groq_api(prompt):
    """Call Groq API for detailed disease analysis and advice"""
    print(f"Calling Groq API with prompt: {prompt[:50]}...")
    
    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "model": GROQ_MODEL,
        "messages": [
            {"role": "system", "content": "You are an expert agricultural advisor specializing in tomato farming and plant diseases."},
            {"role": "user", "content": prompt}
        ],
        "max_tokens": 800,
        "temperature": 0.7
    }
    
    try:
        response = requests.post(
            "https://api.groq.com/openai/v1/chat/completions",
            headers=headers,
            json=payload,
            timeout=30
        )
        
        print(f"Groq API response status: {response.status_code}")
        
        if response.status_code == 200:
            result = response.json()
            if "choices" in result and len(result["choices"]) > 0:
                content = result["choices"][0]["message"]["content"]
                print(f"Groq API response received: {len(content)} characters")
                return content
        
        print(f"Groq API error: {response.status_code} - {response.text}")
        return None
    
    except Exception as e:
        print(f"Error with Groq API: {str(e)}")
        return None

# ===== Fallback to Hugging Face API =====
def call_hf_api(prompt, model_id="mistralai/Mistral-7B-Instruct-v0.2"):
    """Call Hugging Face API as fallback"""
    if not HF_API_TOKEN:
        return None
    
    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
    
    # Format prompt for instruction-tuned models
    formatted_prompt = f"""<s>[INST] {prompt} [/INST]"""
    
    payload = {
        "inputs": formatted_prompt,
        "parameters": {
            "max_new_tokens": 500,
            "temperature": 0.7,
            "top_p": 0.95,
            "do_sample": True
        }
    }
    
    url = f"https://api-inference.huggingface.co/models/{model_id}"
    
    try:
        response = requests.post(url, headers=headers, json=payload, timeout=30)
        
        if response.status_code == 200:
            result = response.json()
            if isinstance(result, list) and len(result) > 0:
                if "generated_text" in result[0]:
                    # Extract just the response part (after the prompt)
                    generated_text = result[0]["generated_text"]
                    # Remove the prompt from the response
                    response_text = generated_text.split("[/INST]")[-1].strip()
                    return response_text
        
        return None
    
    except Exception as e:
        print(f"Error with Hugging Face API: {str(e)}")
        return None

# ===== AI Assistant Functions =====
def ai_assistant_v1(image, prediction, confidence):
    """Use Groq API for Model A versions"""
    if "healthy" in prediction.lower():
        prompt = (
            "You are an agricultural advisor speaking to a farmer. "
            "The tomato crop appears healthy. "
            "Provide detailed preventive tips and best practices for maintaining tomato crop health. "
            "Include information about watering, fertilization, pest prevention, and optimal growing conditions. "
            "Format your response in clear sections with bullet points where appropriate."
        )
    else:
        prompt = (
            f"You are an agricultural advisor speaking to a farmer. "
            f"A disease has been detected in their tomato crop: {prediction} with {confidence:.1f}% confidence. "
            f"Provide detailed advice on how to identify, manage and treat this disease. "
            f"Include information about: "
            f"1) What causes this disease "
            f"2) How it spreads "
            f"3) Specific treatments (both organic and chemical options) "
            f"4) Preventive measures for the future "
            f"Format your response in clear sections with bullet points where appropriate."
        )
    
    # Call Groq API
    response = call_groq_api(prompt)
    
    # If Groq API fails, try Hugging Face API
    if not response:
        response = call_hf_api(prompt)
    
    # If both APIs fail, use fallback information
    if not response:
        # Get fallback information from our database
        info = disease_info.get(prediction, {
            "description": "Information not available for this disease.",
            "causes": "Unknown causes.",
            "treatment": ["Consult with a local agricultural extension service."]
        })
        
        response = f"""
# {prediction}

## Description
{info.get('description', 'No description available.')}

## Causes
{info.get('causes', 'Information not available.')}

## Recommended Treatment
{chr(10).join(f"- {rec}" for rec in info.get('treatment', ['No specific treatment information available.']))}

*Note: This is fallback information. For more detailed advice, please try again later when the AI service is available.*
"""
    
    return response

def ai_assistant_v2(image, prediction, confidence):
    """Use Groq API for Model B versions"""
    if "healthy" in prediction.lower():
        prompt = (
            "You are an agricultural advisor speaking to a farmer. "
            "The tomato crop appears healthy. "
            "Provide detailed preventive tips and best practices for maintaining tomato crop health. "
            "Include information about watering, fertilization, pest prevention, and optimal growing conditions. "
            "Format your response in clear sections with bullet points where appropriate."
        )
    else:
        prompt = (
            f"You are an agricultural advisor speaking to a farmer. "
            f"A disease has been detected in their tomato crop: {prediction} with {confidence:.1f}% confidence. "
            f"Provide detailed advice on how to identify, manage and treat this disease. "
            f"Include information about: "
            f"1) What causes this disease "
            f"2) How it spreads "
            f"3) Specific treatments (both organic and chemical options) "
            f"4) Preventive measures for the future "
            f"Format your response in clear sections with bullet points where appropriate."
        )
    
    # Call Groq API
    response = call_groq_api(prompt)
    
    # If Groq API fails, try Hugging Face API
    if not response:
        response = call_hf_api(prompt)
    
    # If both APIs fail, use fallback information
    if not response:
        # Get fallback information from our database
        info = disease_info.get(prediction, {
            "description": "Information not available for this disease.",
            "causes": "Unknown causes.",
            "treatment": ["Consult with a local agricultural extension service."]
        })
        
        response = f"""
# {prediction}

## Description
{info.get('description', 'No description available.')}

## Causes
{info.get('causes', 'Information not available.')}

## Recommended Treatment
{chr(10).join(f"- {rec}" for rec in info.get('treatment', ['No specific treatment information available.']))}

*Note: This is fallback information. For more detailed advice, please try again later when the AI service is available.*
"""
    
    return response

# ===== Process Function Based on Version =====
def process_version(image, version):
    if image is None:
        return "No image provided."
    
    # --- Version 1.x (Model A) ---
    if version == "1.1":
        result, confidence = predict_model_a(image)
        return f"Model A Prediction: {result} (Confidence: {confidence:.1f}%)\n\nView Model A Training Notebook: https://colab.research.google.com/drive/1FMjs7JmdO6WVoXbzLA-ymwnIKq-GaV6w?usp=sharing"
    
    elif version == "1.2":
        result, confidence = predict_model_a(image)
        advice = ai_assistant_v1(image, result, confidence)
        return f"Model A Prediction: {result} (Confidence: {confidence:.1f}%)\n\nExpert Advice:\n{advice}"
    
    elif version == "1.3":
        cls_result = predict_classifier(image)
        if cls_result != "Tomato Leaf":
            return "Classifier: The image is not a tomato leaf. Please try again with a tomato leaf image."
        
        result, confidence = predict_model_a(image)
        advice = ai_assistant_v1(image, result, confidence)
        return (
            f"Classifier: {cls_result}\n"
            f"Model A Prediction: {result} (Confidence: {confidence:.1f}%)\n\n"
            f"Expert Advice:\n{advice}\n\n"
            f"[View Model A & Classifier Training Notebook](https://colab.research.google.com/drive/1FMjs7JmdO6WVoXbzLA-ymwnIKq-GaV6w?usp=sharing)"
        )
    
    # --- Version 2.x (Model B) ---
    elif version == "2.1":
        result, confidence = predict_model_b(image)
        return f"Model B Prediction: {result} (Confidence: {confidence:.1f}%)\n\n[View Model B Training Notebook](https://colab.research.google.com/drive/1CvoQY40gK2YsMgt4wq9kM2ZSO2c4lzFU?usp=sharing)"
    
    elif version == "2.2":
        result, confidence = predict_model_b(image)
        advice = ai_assistant_v2(image, result, confidence)
        return f"Model B Prediction: {result} (Confidence: {confidence:.1f}%)\n\nExpert Advice:\n{advice}"
    
    elif version == "2.3":
        cls_result = predict_classifier(image)
        if cls_result != "Tomato Leaf":
            return "Classifier: The image is not a tomato leaf. Please try again with a tomato leaf image."
        
        result, confidence = predict_model_b(image)
        advice = ai_assistant_v2(image, result, confidence)
        return (
            f"Classifier: {cls_result}\n"
            f"Model B Prediction: {result} (Confidence: {confidence:.1f}%)\n\n"
            f"Expert Advice:\n{advice}\n\n"
            f"[View Model B & Classifier Training Notebook](https://colab.research.google.com/drive/1CvoQY40gK2YsMgt4wq9kM2ZSO2c4lzFU?usp=sharing)"
        )
    
    else:
        return "Invalid version selected."

# ===== Helper Function to Choose Between Uploaded & Camera Image =====
def combine_images(uploaded, camera):
    return camera if camera is not None else uploaded

# ===== CSS for Theme Switching =====
light_css = """
<style>
body { background-color: white; color: black; }
.gr-button { background-color: #4CAF50; color: white; }
.gr-input, .gr-textbox, .gr-dropdown, .gr-radio, .gr-markdown, .gr-container { background-color: white; color: black; }
</style>
"""

dark_css = """
<style>
body { background-color: #121212 !important; color: #e0e0e0 !important; }
.gr-button { background-color: #555 !important; color: white !important; }
.gr-input, .gr-textbox, .gr-dropdown, .gr-radio, .gr-markdown, .gr-container { background-color: #333 !important; color: #e0e0e0 !important; }
</style>
"""

def update_css(theme):
    if theme == "Dark":
        return dark_css
    else:
        return light_css

# ===== Gradio Interface =====
with gr.Blocks() as demo:
    # Hidden element for CSS injection (initially Light theme)
    css_injector = gr.HTML(update_css("Light"))
    
    gr.Markdown("# 🌿 FarMVi8ioN – AI-powered Crop Monitoring")
    gr.Markdown("Detect tomato leaf diseases and get actionable advice on how to curb them.")
    
    with gr.Row():
        # ----- Left Column (≈30%) -----
        with gr.Column(scale=1):
            version = gr.Dropdown(
                choices=["1.1", "1.2", "1.3", "2.1", "2.2", "2.3"],
                label="Select Version",
                value="1.3",
                info="Versions 1.x use Model A; Versions 2.x use Model B."
            )
            
            theme_choice = gr.Radio(
                choices=["Light", "Dark"],
                label="Select Theme",
                value="Light"
            )
            
            gr.Markdown("### Notebook Links")
            gr.Markdown(
                """
                **For Model A:**
                - Model A Only: [Training Notebook](https://colab.research.google.com/drive/1FMjs7JmdO6WVoXbzLA-ymwnIKq-GaV6w?usp=sharing)
                - Model A & Classifier: [Training Notebook](https://colab.research.google.com/drive/1CvoQY40gK2YsMgt4wq9kM2ZSO2c4lzFU?usp=sharing)
                
                **For Model B:**
                - Model B Only: [Training Notebook](https://colab.research.google.com/drive/1CvoQY40gK2YsMgt4wq9kM2ZSO2c4lzFU?usp=sharing)
                - Model B & Classifier: [Training Notebook](https://colab.research.google.com/drive/1CvoQY40gK2YsMgt4wq9kM2ZSO2c4lzFU?usp=sharing)
                """
            )
        
        # ----- Right Column (≈70%) -----
        with gr.Column(scale=2):
            image_input = gr.Image(label="📂 Upload Tomato Leaf Image", type="pil", sources=["upload", "webcam", "clipboard"])
            submit = gr.Button("🔍 Analyze", variant="primary")
    
    output = gr.Markdown(label="📝 Diagnosis & Advice")
    
    # Update CSS dynamically based on theme selection
    theme_choice.change(fn=update_css, inputs=theme_choice, outputs=css_injector)
    
    # When submit is clicked, combine image inputs and process the selected version
    submit.click(
        fn=lambda img, ver: process_version(img, ver),
        inputs=[image_input, version],
        outputs=output
    )

demo.launch()