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

def stream_response(container, response):
    """Gère le streaming de la réponse avec un affichage progressif"""
    mode = 'starting'
    thinking_placeholder = None
    answer_placeholder = None
    thinking_text = ""
    answer_text = ""
    
    for chunk in response:
        for part in chunk.candidates[0].content.parts:
            if part.thought:
                if mode != "thinking":
                    if thinking_placeholder is None:
                        with container.expander("Voir le raisonnement", expanded=False):
                            thinking_placeholder = st.empty()
                    mode = "thinking"
                thinking_text += part.text
                thinking_placeholder.markdown(thinking_text)
            else:
                if mode != "answering":
                    if answer_placeholder is None:
                        answer_placeholder = container.empty()
                        container.subheader("Réponse")
                    mode = "answering"
                answer_text += part.text
                answer_placeholder.markdown(answer_text)

def main():
    st.title("Analyseur d'Images Géométriques avec Gemini")
    
    # Utilisation des secrets pour la clé API
    try:
        api_key = st.secrets["GEMINI_API_KEY"]
    except Exception as e:
        st.error("Erreur: dans les secrets.")
        
        return
    
    # Initialisation du client
    try:
        client = genai.Client(
            api_key=api_key,
            http_options={'api_version':'v1alpha'}
        )
    except Exception as e:
        st.error(f"Erreur lors de l'initialisation du client: {e}")
        return
    
    # Upload de l'image
    uploaded_file = st.file_uploader("Choisissez une image géométrique", type=['png', 'jpg', 'jpeg'])
    
    if uploaded_file:
        try:
            # Affichage de l'image
            image = Image.open(uploaded_file)
            st.image(image, caption="Image téléchargée", use_column_width=True)
            
            # Configuration du modèle
            model_name = "gemini-2.0-flash-thinking-exp-01-21"
            
            if st.button("Analyser l'image"):
                # Création d'un conteneur pour la réponse en streaming
                response_container = st.container()
                
                with st.spinner("Analyse en cours..."):
                    try:
                        # Génération de la réponse en streaming
                        response = client.models.generate_content_stream(
                            model=model_name,
                            config={'thinking_config': {'include_thoughts': True}},
                            contents=[
                                image,
                                "What's the area of the overlapping region?"
                            ]
                        )
                        
                        # Streaming de la réponse
                        stream_response(response_container, response)
                                
                    except Exception as e:
                        st.error(f"Erreur lors de l'analyse: {e}")
                        
        except Exception as e:
            st.error(f"Erreur lors du traitement de l'image: {e}")

if __name__ == "__main__":
    main()