import streamlit as st
from google import genai
from PIL import Image
import os
import tempfile
import subprocess

# Configuration de l'API Gemini
gen = os.environ['GOOGLE_API_KEY']
client = genai.Client(api_key=gen)

def ensure_latex_packages():
    """Vérifie si les packages LaTeX nécessaires sont installés."""
    try:
        subprocess.run(['pdflatex', '--version'], capture_output=True, check=True)
        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        st.error("pdflatex n'est pas installé. Veuillez installer TeX Live ou MiKTeX.")
        return False

def generate_latex_response(image_path, question):
    """Génère une réponse en format LaTeX depuis le modèle Gemini."""
    try:
        # Charge l'image pour l'API
        with open(image_path, 'rb') as img_file:
            image_data = img_file.read()
            
        # Prépare le template LaTeX
        latex_template = r"""\documentclass{article}
\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{amsmath}
\usepackage{amsfonts}
\usepackage{amssymb}
\begin{document}
%s
\end{document}
"""
        
        # Génère la réponse avec Gemini
        response = client.models.generate_content(
            model="gemini-2.0-flash-thinking-exp",
            contents=[{"mime_type": "image/jpeg", "data": image_data}, question]
        )
        
        # Extrait la réponse
        answer = response.text
        
        # Formate la réponse en LaTeX complet
        return latex_template % answer
        
    except Exception as e:
        st.error(f"Erreur lors de la génération de la réponse: {str(e)}")
        return None

def latex_to_pdf(latex_content):
    """Convertit le contenu LaTeX en PDF."""
    with tempfile.TemporaryDirectory() as temp_dir:
        # Crée le fichier tex
        tex_file = os.path.join(temp_dir, "output.tex")
        with open(tex_file, "w", encoding="utf-8") as f:
            f.write(latex_content)
        
        try:
            # Exécute pdflatex
            process = subprocess.run(
                ["pdflatex", "-interaction=nonstopmode", tex_file],
                cwd=temp_dir,
                capture_output=True,
                text=True,
                check=True
            )
            
            # Vérifie si le PDF a été généré
            pdf_path = os.path.join(temp_dir, "output.pdf")
            if os.path.exists(pdf_path):
                with open(pdf_path, "rb") as f:
                    return f.read()
            else:
                st.error("Le fichier PDF n'a pas été généré")
                st.code(process.stdout)
                return None
                
        except subprocess.CalledProcessError as e:
            st.error(f"Erreur lors de la compilation LaTeX: {str(e)}")
            st.code(e.stdout)
            return None

def main():
    st.title("Application Gemini 2.0 avec Export PDF")
    
    if not ensure_latex_packages():
        return
    
    uploaded_file = st.file_uploader("Choisissez une image...", type=["png", "jpg", "jpeg"])
    
    if uploaded_file:
        # Sauvegarde temporaire de l'image
        with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as temp_file:
            temp_file.write(uploaded_file.getvalue())
            temp_path = temp_file.name
            
        st.image(uploaded_file, caption="Image téléchargée")
        
        question = st.text_input("Entrez votre question:")
        
        if st.button("Générer la réponse") and question:
            with st.spinner("Génération de la réponse..."):
                # Génère la réponse LaTeX
                latex_content = generate_latex_response(temp_path, question)
                
                if latex_content:
                    st.markdown("### Réponse LaTeX générée:")
                    st.text(latex_content)
                    
                    # Convertit en PDF
                    with st.spinner("Conversion en PDF..."):
                        pdf_data = latex_to_pdf(latex_content)
                        
                        if pdf_data:
                            st.success("PDF généré avec succès!")
                            st.download_button(
                                "Télécharger le PDF",
                                pdf_data,
                                file_name="reponse.pdf",
                                mime="application/pdf"
                            )
        
        # Nettoie le fichier temporaire
        os.unlink(temp_path)

if __name__ == "__main__":
    main()