import streamlit as st
import os
import base64
from google import genai

# Initialize Gemini Client with API key from environment
client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))

# Helper to encode image for Gemini
def process_image(image_file):
    image_bytes = image_file.read()
    encoded_image = base64.b64encode(image_bytes).decode("utf-8")
    return {
        "inline_data": {
            "mime_type": image_file.type,
            "data": encoded_image
        }
    }

# Main App
def main():
    st.set_page_config(page_title="🌿 Leaf Disease Detector", layout="centered")
    st.title("🌱 Leaf Disease Detector")
    st.write("Upload an image of a plant leaf, and we'll analyze it to detect possible diseases, "
             "provide treatment suggestions, and find real-world statistics from the internet about this disease.")

    uploaded_image = st.file_uploader("Upload a leaf image (JPG or PNG)", type=["jpg", "jpeg", "png"])

    if uploaded_image:
        st.image(uploaded_image, caption="Uploaded Leaf", use_container_width=True)
        image_data = process_image(uploaded_image)

        if st.button("🧪 Detect Disease"):
            with st.spinner("Analyzing leaf and searching for information..."):
                content_blocks = [
                    {
                        "text": (
                            "You are a plant disease diagnostic AI. Analyze the uploaded leaf image and do the following:\n\n"
                            "1. Identify any plant disease visible on the leaf.\n"
                            "2. Provide the disease name and a short scientific description.\n"
                            "3. Suggest treatment or prevention methods farmers can use.\n"
                            "4. Use **Google Search** to find:\n"
                            "   - Estimated global or regional financial losses due to this disease\n"
                            "   - Impactful statistics or facts due to this disease (e.g., crops affected, common locations, yield reduction, etc.)\n"
                            "If the leaf looks healthy, clearly state that."
                        )
                    },
                    image_data
                ]

                response = client.models.generate_content(
                    model="gemini-2.5-flash",
                    contents=content_blocks,
                    config={"tools": [{"google_search": {}}]}
                )

                st.success("📋 Analysis Result:")
                st.markdown(response.text)

if __name__ == "__main__":
    main()