import streamlit as st
import requests
from PIL import Image
from io import BytesIO
import os
import google.generativeai as genai

# Obtener la clave API de las variables de entorno
gemini_api_key = os.getenv("GEMINI_API_KEY")
genai.configure(api_key=gemini_api_key)
model = genai.GenerativeModel('gemini-pro')

# Configuración de la API de predicción de enfermedades
API_URL = "https://api-inference.huggingface.co/models/riabayonaor/modelo_prediccion_enfermedades_pepinos"

def query(image_bytes):
    response = requests.post(API_URL, data=image_bytes)
    if response.status_code != 200:
        st.error(f"Error al clasificar la imagen: {response.status_code}")
        return None
    return response.json()

def gemini_query(prompt):
    response = model.generate_content(prompt)
    return response.text

def main():
    st.set_page_config(page_title="Predicción de Enfermedades en Pepinos")
    st.title("Predicción de Enfermedades en Pepinos")
    st.write("Sube una foto de una planta de pepino o un pepino para clasificar posibles enfermedades y obtener soluciones.")

    uploaded_file = st.file_uploader("Sube una foto de una planta de pepino o un pepino", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption='Imagen subida.', use_column_width=True)
        st.write("Clasificando...")

        # Convertir la imagen a bytes
        img_byte_arr = BytesIO()
        image.save(img_byte_arr, format='PNG')
        img_byte_arr = img_byte_arr.getvalue()

        # Enviar la imagen al modelo de Hugging Face
        predictions = query(img_byte_arr)

        if predictions is not None:
            if "error" not in predictions:
                # Suponiendo que las predicciones están en el formato [{label: "Downy Mildew", score: 0.95}, {label: "Fresh Cucumber", score: 0.05}]
                top_prediction = max(predictions, key=lambda x: x["score"])
                class_label = top_prediction['label']
                st.write(f"Predicción principal: {class_label} con confianza {top_prediction['score']:.2f}")

                # Usar la etiqueta principal para el modelo Gemini
                prompt = f"Esta enfermedad es {class_label}. Explica qué es y sugiere posibles insecticidas o soluciones en español."

                # Llamar al modelo Gemini
                explanation = gemini_query(prompt)
                
                st.write(f"Esta enfermedad es {class_label}:")
                st.write(explanation)
            else:
                st.write("No se pudo clasificar la imagen.")
        else:
            st.write("No se pudo clasificar la imagen.")

if __name__ == "__main__":
    main()