import streamlit as st
from google import genai
from google.genai import types
from PIL import Image
import json
import logging
from typing import Optional, Generator, Any
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 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
                ]
            )
            return response
        except Exception as e:
            logger.error(f"Erreur lors de l'analyse de l'image: {e}")
            raise

def stream_response(container, response: Generator) -> None:
    """Gère le streaming de la réponse avec un affichage progressif et une gestion d'erreurs robuste"""
    mode = 'starting'
    thinking_placeholder = None
    answer_placeholder = None
    thinking_text = ""
    answer_text = ""
    
    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
                        
                    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 += 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 += 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")
    
    # 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. La réponse doit être bien présentée et espacée pour faciliter la lecture. Réponds en français."
            
            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()