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

# Configuration de la page
st.set_page_config(page_title="Résolveur d'exercices")

# Configuration de l'API Gemini
GOOGLE_API_KEY = "AIzaSyC_zxN9IHjEAxIoshWPzMfgb9qwMsu5t5Y"
client = genai.Client(
    api_key=GOOGLE_API_KEY,
    http_options={'api_version': 'v1alpha'},
)

def main():
    st.title("Résolveur d'exercices")
    
    # Zone de téléchargement d'image
    uploaded_file = st.file_uploader("Téléchargez votre exercice", type=["png", "jpg", "jpeg"])
    
    if uploaded_file:
        # Affichage de l'image
        image = Image.open(uploaded_file)
        st.image(image, caption="Image téléchargée", use_column_width=True)
        
        # Bouton pour lancer l'analyse
        if st.button("Résoudre l'exercice"):
            # Conversion de l'image en base64
            buffered = io.BytesIO()
            image.save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode()
            
            # Création de deux colonnes : une pour les pensées, une pour la réponse
            thoughts_col, answer_col = st.columns([1, 2])
            
            with thoughts_col:
                # Création d'un expander pour les pensées
                with st.expander("Pensées du modèle", expanded=True):
                    thoughts_placeholder = st.empty()
            
            with answer_col:
                # Placeholder pour la réponse
                answer_placeholder = st.empty()
            
            try:
                # Génération du contenu avec streaming
                response = 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}},
                        "Résous cet exercice. ça doit être bien présentable et espacé afin d'être facile à lire."
                    ]
                )
                
                # Variables pour accumuler le contenu
                thoughts = []
                answer = []
                
                # Traitement du stream
                for chunk in response:
                    for part in chunk.candidates[0].content.parts:
                        if part.thought:
                            thoughts.append(part.text)
                            thoughts_placeholder.markdown("\n\n".join(thoughts))
                        else:
                            answer.append(part.text)
                            answer_placeholder.markdown("\n\n".join(answer))
                
            except Exception as e:
                st.error(f"Une erreur s'est produite : {str(e)}")

if __name__ == "__main__":
    main()