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

# Configuration du logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

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 = ""
    
    try:
        for chunk in response:
            logger.debug(f"Chunk reçu: {chunk}")
            
            try:
                for part in chunk.candidates[0].content.parts:
                    if hasattr(part, 'thought') and 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)
            except Exception as e:
                logger.error(f"Erreur lors du traitement d'un chunk: {e}")
                logger.debug(f"Contenu du chunk problématique: {chunk}")
                raise
                
    except Exception as e:
        logger.error(f"Erreur dans le streaming de la réponse: {e}")
        raise

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:
        logger.error(f"Erreur dans la récupération des secrets: {e}")
        st.error("Erreur: Impossible d'accéder aux secrets de l'application.")
        return
    
    # Initialisation du client
    try:
        client = genai.Client(
            api_key=api_key,
            http_options={'api_version':'v1alpha'}
        )
    except Exception as e:
        logger.error(f"Erreur d'initialisation du client: {e}")
        st.error(f"Erreur lors de l'initialisation du client Gemini: {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_container_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:
                        # Préparation de la requête
                        request = {
                            "contents": [
                                {"image": image},
                                {"text": "What's the area of the overlapping region?"}
                            ]
                        }
                        
                        logger.debug(f"Envoi de la requête au modèle: {request}")
                        
                        # 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:
                        logger.error(f"Erreur lors de l'analyse: {e}", exc_info=True)
                        st.error(f"Erreur lors de l'analyse de l'image: {str(e)}")
                        
        except Exception as e:
            logger.error(f"Erreur lors du traitement de l'image: {e}", exc_info=True)
            st.error(f"Erreur lors du traitement de l'image: {str(e)}")

if __name__ == "__main__":
    main()