import streamlit as st
from google import genai
from google.genai.types import GenerateContentConfig, ThinkingConfig

def main():
    st.title("Interface Gemini API avec Thoughts")
    
    # Zone de saisie pour la clé API
    api_key = st.text_input("Entrez votre clé API Google", type="password")
    
    # Zone de texte pour le prompt
    prompt = st.text_area("Entrez votre prompt")
    
    if st.button("Générer"):
        if not api_key:
            st.error("Veuillez entrer une clé API")
            return
        
        if not prompt:
            st.warning("Veuillez entrer un prompt")
            return
            
        try:
            # Configuration du client
            client = genai.Client(
                api_key=api_key,
                http_options={'api_version': 'v1alpha'}
            )
            
            # Configuration pour inclure les chain-of-thoughts
            config = GenerateContentConfig(
                thinking_config=ThinkingConfig(include_thoughts=True)
            )
            
            # Affichage d'un spinner pendant la génération
            with st.spinner("Génération en cours..."):
                response = client.models.generate_content(
                    model="gemini-2.0-flash-thinking-exp-01-21",
                    contents=prompt,
                    config=config,
                )
            
            st.subheader("Résultats")
            # Affichage dans la console pour debug
            print(response)
            print(response.candidates[0].content.parts)
            # Parcourir les parties de la réponse
            for part in response.candidates[0].content.parts:
                # Si le modèle retourne un champ 'thought', on l'affiche dans un expander dédié
                if part.thought:
                    with st.expander("💭 Thoughts", expanded=True):
                        st.write(part.text)
                else:
                    with st.expander("📝 Réponse", expanded=True):
                        st.write(part.text)
                        
        except Exception as e:
            st.error(f"Une erreur s'est produite: {str(e)}")

if __name__ == "__main__":
    # Configuration de la page Streamlit
    st.set_page_config(
        page_title="Interface Gemini API",
        page_icon="🤖",
        layout="wide"
    )
    main()