import streamlit as st
from google import genai
from google.genai import types
import os
from PIL import Image
import io
import base64
import json

GOOGLE_API_KEY =  os.environ.get("GEMINI_API_KEY")

client = genai.Client(
    api_key=GOOGLE_API_KEY,
    http_options={'api_version': 'v1alpha'},
)

st.title("Résolution d'exercice avec Gemini")

uploaded_file = st.file_uploader("Télécharger une image de l'exercice", type=["png", "jpg", "jpeg"])

if uploaded_file is not None:
    image_data = uploaded_file.read()
    img = Image.open(io.BytesIO(image_data))

    buffered = io.BytesIO()
    img.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()

    if st.button("Résoudre l'exercice"):
        response_area = st.empty()  # Placeholder pour la réponse
        mode_area = st.empty()      # Placeholder pour le mode (thinking/answering)

        mode = 'starting'
        try:
            responses = client.models.generate_content_stream(
                    model="gemini-2.0-flash-thinking-exp-01-21",
                    config={'thinking_config': {'include_thoughts': True}},
                    contents=[
                        {'inline_data': {'mime_type': 'image/png', 'data': img_str}},
                        "Resous cette exercice. ça doit être bien présentable et espacé afin d'être facile à lire."
                    ]
            )

            full_response = "" # Pour accumuler toute la réponse pour un affichage final propre

            for response in responses:
                for part in response.parts:
                    if hasattr(part, 'thought') and part.thought: # Vérifier si la partie a un attribut 'thought' et s'il est True
                        if mode != "thinking":
                            mode_area.markdown("**Gemini réfléchit... 🤔**") # Afficher "thinking" dans le placeholder
                            mode = "thinking"
                    elif hasattr(part, 'text'): # Vérifier si la partie a un attribut 'text' (pour la réponse)
                        if mode != "answering":
                            mode_area.markdown("**Réponse de Gemini :**") # Afficher "answering" dans le placeholder
                            mode = "answering"

                        text_chunk = part.text or "" # Gérer le cas où part.text pourrait être None
                        full_response += text_chunk
                        response_area.markdown(full_response) # Mettre à jour la réponse en continu

            mode_area.empty() # Effacer le message "thinking/answering" une fois terminé
            st.success("Exercice résolu !") # Message de succès

        except Exception as e:
            print(f"Erreur lors de la génération: {e}")
            mode_area.empty() # Effacer les messages de mode en cas d'erreur
            response_area.error(f"Erreur: {e}") # Afficher l'erreur dans le placeholder de réponse