import streamlit as st
from google import genai
from google.genai import types
from PIL import Image
import json
import logging
import re
from typing import Optional, Generator, Any, Tuple
import sys
from pathlib import Path

# Configuration du logging
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler(Path('app.log'))
    ]
)
logger = logging.getLogger(__name__)

class LatexFormatter:
    """Classe pour gérer le formatage LaTeX"""
    
    @staticmethod
    def format_inline_latex(text: str) -> str:
        """Formate le LaTeX inline en ajoutant les délimiteurs appropriés"""
        # Remplace les délimiteurs $ simples par $$ pour inline math
        text = re.sub(r'(?<!\\)\$([^$]+)(?<!\\)\$', r'$\1$', text)
        return text
    
    @staticmethod
    def format_block_latex(text: str) -> str:
        """Formate les blocs LaTeX"""
        # Ajoute des sauts de ligne avant et après les blocs LaTeX
        text = re.sub(r'\$\$(.*?)\$\$', r'\n\n$$\1$$\n\n', text, flags=re.DOTALL)
        return text
    
    @staticmethod
    def enhance_latex_display(text: str) -> str:
        """Améliore l'affichage LaTeX global"""
        # Gère les cas spéciaux de LaTeX
        text = text.replace('\\[', '$$').replace('\\]', '$$')  # Convertit \[ \] en $$
        text = text.replace('\\(', '$').replace('\\)', '$')    # Convertit \( \) en $
        
        # Assure l'espacement correct autour des équations
        text = re.sub(r'(\$\$.*?\$\$)', r'\n\n\1\n\n', text, flags=re.DOTALL)
        
        # Améliore l'alignement des environnements mathématiques
        text = re.sub(r'(\\begin{[^}]*})(.*?)(\\end{[^}]*})', 
                     r'\n\n$$\1\2\3$$\n\n', 
                     text, flags=re.DOTALL)
        
        return text

class GeminiClient:
    """Classe pour gérer les interactions avec l'API Gemini"""
    def __init__(self, api_key: str):
        self.client = None
        self.init_client(api_key)
        
    def init_client(self, api_key: str) -> None:
        """Initialise le client Gemini"""
        try:
            self.client = genai.Client(
                api_key=api_key,
                http_options={'api_version': 'v1alpha'}
            )
        except Exception as e:
            logger.error(f"Erreur d'initialisation du client Gemini: {e}")
            raise RuntimeError(f"Impossible d'initialiser le client Gemini: {e}")

    def analyze_image(self, image: Image.Image, prompt: str, model_name: str) -> Generator:
        """Analyse une image avec Gemini"""
        if not self.client:
            raise RuntimeError("Client Gemini non initialisé")
            
        try:
            response = self.client.models.generate_content_stream(
                model=model_name,
                config={'thinking_config': {'include_thoughts': True}},
                contents=[
                    image,
                    prompt + " Si ta réponse contient des expressions mathématiques, utilise la notation LaTeX avec les délimiteurs appropriés ($ pour inline, $$ pour les blocs)."
                ]
            )
            return response
        except Exception as e:
            logger.error(f"Erreur lors de l'analyse de l'image: {e}")
            raise

def format_text_with_latex(text: str) -> str:
    """Formate le texte en appliquant les améliorations LaTeX"""
    formatter = LatexFormatter()
    text = formatter.format_inline_latex(text)
    text = formatter.format_block_latex(text)
    text = formatter.enhance_latex_display(text)
    return text

def stream_response(container, response: Generator) -> None:
    """Gère le streaming de la réponse avec support LaTeX amélioré"""
    mode = 'starting'
    thinking_placeholder = None
    answer_placeholder = None
    thinking_text = ""
    answer_text = ""
    
    # Configuration CSS pour améliorer l'affichage LaTeX
    st.markdown("""
        <style>
        .katex { font-size: 1.1em !important; }
        .katex-display { overflow: auto hidden; }
        .markdown-text-container { max-width: 100%; }
        </style>
        """, unsafe_allow_html=True)
    
    try:
        for chunk in response:
            logger.debug(f"Chunk reçu: {chunk}")
            
            if not isinstance(chunk, (dict, types.GenerateContentResponse)):
                logger.warning(f"Format de chunk invalide reçu: {type(chunk)}")
                continue
                
            try:
                candidates = getattr(chunk, 'candidates', None)
                if not candidates or not len(candidates):
                    logger.warning("Pas de candidats dans le chunk")
                    continue
                    
                content = getattr(candidates[0], 'content', None)
                if not content:
                    logger.warning("Pas de contenu dans le premier candidat")
                    continue
                    
                parts = getattr(content, 'parts', [])
                for part in parts:
                    has_thought = False
                    try:
                        has_thought = hasattr(part, 'thought') and part.thought
                    except Exception as e:
                        logger.warning(f"Erreur lors de la vérification de thought: {e}")
                    
                    text = getattr(part, 'text', '')
                    if not text:
                        continue
                    
                    # Formatage LaTeX du texte
                    formatted_text = format_text_with_latex(text)
                        
                    if has_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 += formatted_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 += formatted_text
                        answer_placeholder.markdown(answer_text)
                        
            except json.JSONDecodeError as e:
                logger.error(f"Erreur de décodage JSON: {e}")
                continue
            except Exception as e:
                logger.error(f"Erreur lors du traitement d'un chunk: {e}")
                logger.debug(f"Contenu du chunk problématique: {chunk}")
                continue
                
    except Exception as e:
        logger.error(f"Erreur fatale dans le streaming de la réponse: {e}")
        if not answer_text and not thinking_text:
            container.error("Une erreur est survenue lors de l'analyse de l'image. Veuillez réessayer.")
        raise
    finally:
        if not answer_text and not thinking_text:
            container.warning("Aucune réponse n'a pu être générée. Veuillez réessayer.")

def validate_image(uploaded_file) -> Optional[Image.Image]:
    """Valide et ouvre une image téléchargée"""
    try:
        image = Image.open(uploaded_file)
        return image
    except Exception as e:
        logger.error(f"Erreur lors de l'ouverture de l'image: {e}")
        st.error("L'image n'a pas pu être ouverte. Veuillez vérifier le format.")
        return None

def main():
    st.set_page_config(
        page_title="Mariam M-0",
        page_icon="🔍",
        layout="wide"
    )
    
    st.title("Mariam M-0")
    
    # Configuration supplémentaire pour LaTeX
    st.markdown("""
        <style>
        /* Styles globaux pour LaTeX */
        .latex-container { 
            margin: 1em 0;
            padding: 0.5em;
            background: #f9f9f9;
            border-radius: 4px;
        }
        /* Amélioration de l'espacement des équations */
        .katex-display { 
            margin: 1.5em 0 !important; 
            padding: 0.5em 0;
        }
        /* Amélioration du rendu sur mobile */
        @media (max-width: 768px) {
            .katex { font-size: 1em !important; }
        }
        </style>
    """, unsafe_allow_html=True)
    
    # Récupération de 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:
        gemini_client = GeminiClient(api_key)
    except Exception as e:
        st.error(f"Erreur lors de l'initialisation du client Gemini: {e}")
        return
    
    # Interface utilisateur
    uploaded_file = st.file_uploader(
        "Choisissez une image géométrique",
        type=['png', 'jpg', 'jpeg'],
        help="Formats supportés: PNG, JPG, JPEG"
    )
    
    if uploaded_file:
        image = validate_image(uploaded_file)
        if image:
            st.image(image, caption="Image téléchargée", use_container_width=True)
            
            model_name = "gemini-2.0-flash-thinking-exp-01-21"
            prompt = "Résous cet exercice mathématique. La réponse doit être bien présentée et espacée pour faciliter la lecture. Réponds en français et utilise la notation LaTeX pour toutes les expressions mathématiques."
            
            if st.button("Analyser l'image", type="primary"):
                response_container = st.container()
                
                with st.spinner("Analyse en cours..."):
                    try:
                        response = gemini_client.analyze_image(image, prompt, model_name)
                        stream_response(response_container, response)
                    except Exception as e:
                        logger.error(f"Erreur lors de l'analyse: {e}", exc_info=True)
                        st.error("Une erreur est survenue lors de l'analyse. Veuillez réessayer.")

if __name__ == "__main__":
    main()