import streamlit as st
import requests
from PIL import Image
from io import BytesIO
from dotenv import load_dotenv
import os

# Cargar variables de entorno desde el archivo .env
load_dotenv()
huggingface_api_key = os.getenv("HUGGINGFACE_API_KEY")

# Configuración de la API
API_URL = "https://api-inference.huggingface.co/models/riabayonaor/modelo_prediccion_enfermedades_pepinos"
META_LLAMA_API_URL = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-8B-Instruct"
headers = {"Authorization": f"Bearer {huggingface_api_key}"}

# Lista de clases del modelo
classes = ["Anthracnose", "Bacterial Wilt", "Belly Rot", "Downy Mildew",
           "Fresh Cucumber", "Fresh Leaf", "Gummy Stem Blight", "Pythium Fruit Rot"]

def query(image_bytes):
    response = requests.post(API_URL, headers=headers, 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 llama_query(prompt):
    response = requests.post(META_LLAMA_API_URL, headers=headers, json={"inputs": prompt})
    if response.status_code != 200:
        st.error(f"Error al obtener respuesta del modelo Meta Llama: {response.status_code}")
        return None
    return response.json()

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: "0", score: 0.95}, {label: "1", score: 0.05}]
                top_prediction = max(predictions, key=lambda x: x["score"])
                class_label = classes[int(top_prediction['label'])]
                st.write(f"Predicción principal: {class_label} con confianza {top_prediction['score']:.2f}")

                # Usar la etiqueta principal para el modelo de Meta Llama
                prompt = f"Esta enfermedad es {class_label}. Explica qué es y sugiere posibles insecticidas o soluciones."

                # Llamar al modelo Meta Llama
                llama_response = llama_query(prompt)
                
                if llama_response is not None:
                    if "error" not in llama_response:
                        explanation = llama_response[0]["generated_text"]
                        st.write("Explicación y posibles soluciones:")
                        st.write(explanation)
                    else:
                        st.write("No se pudo obtener una respuesta del modelo Meta Llama.")
                else:
                    st.write("No se pudo obtener una respuesta del modelo Meta Llama.")
            else:
                st.write("No se pudo clasificar la imagen.")
        else:
            st.write("No se pudo clasificar la imagen.")

if __name__ == "__main__":
    main()