import gradio as gr
import base64
from io import BytesIO
from PIL import Image
import pandas as pd
import traceback
import threading
import queue
import time

from rag_chatbot import MultimodalRAGChatbot
from langchain.callbacks.base import BaseCallbackHandler

# Handler personnalisé pour capturer les tokens en streaming
class GradioStreamingHandler(BaseCallbackHandler):
    def __init__(self):
        self.tokens_queue = queue.Queue()
        self.full_text = ""
        
    def on_llm_new_token(self, token, **kwargs):
        self.tokens_queue.put(token)
        self.full_text += token

# Fonction pour créer un objet Image à partir des données base64
def base64_to_image(base64_data):
    """Convertit une image base64 en objet Image pour l'affichage direct"""
    try:
        if not base64_data:
            return None
        image_bytes = base64.b64decode(base64_data)
        image = Image.open(BytesIO(image_bytes))
        return image
    except Exception as e:
        print(f"Erreur lors de la conversion d'image: {e}")
        return None

# Configuration pour initialiser le chatbot
qdrant_url = "http://localhost:6333"
qdrant_collection_name = "my_documents"
embedding_model = "mxbai-embed-large"
ollama_url = "http://127.0.0.1:11434"
default_model = "llama3.1"

# Liste des modèles disponibles
AVAILABLE_MODELS = ["llama3.1", "llama3.2", "deepseek-r1:14b"]

# Initialiser le chatbot RAG avec le modèle par défaut
rag_bot = MultimodalRAGChatbot(
    qdrant_url=qdrant_url,
    qdrant_collection_name=qdrant_collection_name,
    ollama_model=default_model,
    embedding_model=embedding_model,
    ollama_url=ollama_url
)
print(f"Chatbot initialisé avec modèle: {default_model}")

# Variables globales pour stocker les images et tableaux de la dernière requête
current_images = []
current_tables = []

# Fonction pour changer de modèle
def change_model(model_name):
    global rag_bot
    
    try:
        # Réinitialiser le chatbot avec le nouveau modèle
        rag_bot = MultimodalRAGChatbot(
            qdrant_url=qdrant_url,
            qdrant_collection_name=qdrant_collection_name,
            ollama_model=model_name,
            embedding_model=embedding_model,
            ollama_url=ollama_url
        )
        print(f"Modèle changé pour: {model_name}")
        return f"✅ Modèle changé pour: {model_name}"
    except Exception as e:
        print(f"Erreur lors du changement de modèle: {e}")
        return f"❌ Erreur: {str(e)}"

# Fonction de traitement des requêtes avec support du streaming dans Gradio
def process_query(message, history, streaming, show_sources, max_images):
    global current_images, current_tables
    
    if not message.strip():
        return history, "", None, None
    
    current_images = []
    current_tables = []
    
    try:
        if streaming:
            # Version avec streaming dans Gradio
            history = history + [(message, "")]
            
            # 1. Récupérer les documents pertinents
            docs = rag_bot._retrieve_relevant_documents(message)
            
            # 2. Préparer le contexte et l'historique
            context = rag_bot._format_documents(docs)
            history_text = rag_bot._format_chat_history()
            
            # 3. Préparer le prompt
            from langchain.prompts import ChatPromptTemplate
            prompt_template = ChatPromptTemplate.from_template("""
            Tu es un assistant documentaire spécialisé qui utilise toutes les informations disponibles dans le contexte fourni.

            Instructions spécifiques:
            1. Pour chaque image mentionnée dans le contexte, inclue TOUJOURS dans ta réponse:
            - La légende/caption exacte de l'image
            - La source et le numéro de page
            - Une description brève de ce qu'elle montre

            2. Pour chaque tableau mentionné dans le contexte, inclue TOUJOURS:
            - Le titre/caption exact du tableau
            - La source et le numéro de page
            - Ce que contient et signifie le tableau

            3. Lorsque tu cites des équations mathématiques:
            - Utilise la syntaxe LaTeX exacte comme dans le document ($...$ ou $$...$$)
            - Reproduis-les fidèlement sans modification

            4. IMPORTANT: Ne pas inventer d'informations - si une donnée n'est pas explicitement fournie dans le contexte,
            indique clairement "Cette information n'est pas disponible dans les documents fournis."

            5. Cite précisément les sources pour chaque élément d'information (format: [Source, Page]).

            Historique de conversation:
            {chat_history}

            Contexte (à utiliser pour répondre):
            {context}

            Question: {question}

            Réponds de façon structurée et précise en intégrant activement les images, tableaux et équations disponibles dans le contexte.
            """)
            
            # 4. Formater les messages pour le LLM
            messages = prompt_template.format_messages(
                chat_history=history_text,
                context=context,
                question=message
            )
            
            # 5. Créer un handler de streaming personnalisé
            from langchain_ollama import ChatOllama
            handler = GradioStreamingHandler()
            
            # 6. Créer un modèle LLM avec notre handler
            streaming_llm = ChatOllama(
                model=rag_bot.llm.model,
                base_url=rag_bot.llm.base_url,
                streaming=True,
                callbacks=[handler]
            )
            
            # 7. Lancer la génération dans un thread pour ne pas bloquer l'UI
            def generate_response():
                streaming_llm.invoke(messages)
            
            thread = threading.Thread(target=generate_response)
            thread.start()
            
            # 8. Récupérer les tokens et mettre à jour l'interface
            partial_response = ""
            
            # Attendre les tokens avec un timeout
            while thread.is_alive() or not handler.tokens_queue.empty():
                try:
                    token = handler.tokens_queue.get(timeout=0.05)
                    partial_response += token
                    history[-1] = (message, partial_response)
                    yield history, "", None, None
                except queue.Empty:
                    continue
            
            # 9. Thread terminé, mettre à jour l'historique de conversation du chatbot
            rag_bot.chat_history.append({"role": "user", "content": message})
            rag_bot.chat_history.append({"role": "assistant", "content": partial_response})
            
            # 10. Récupérer les sources, images, tableaux
            texts, images, tables = rag_bot._process_documents(docs)
            
            # Préparer les informations sur les sources
            source_info = ""
            if texts:
                source_info += f"📚 {len(texts)} textes • "
            if images:
                source_info += f"🖼️ {len(images)} images • "
            if tables:
                source_info += f"📊 {len(tables)} tableaux"
            
            if source_info:
                source_info = "Sources trouvées: " + source_info
            
            # 11. Traiter les images
            if show_sources and images:
                images = images[:max_images]
                for img in images:
                    img_data = img.get("image_data")
                    if img_data:
                        image = base64_to_image(img_data)
                        if image:
                            current_images.append({
                                "image": image,
                                "caption": img.get("caption", ""),
                                "source": img.get("source", ""),
                                "page": img.get("page", ""),
                                "description": img.get("description", "")
                            })
            
            # 12. Traiter les tableaux
            if show_sources and tables:
                for table in tables:
                    current_tables.append({
                        "data": rag_bot.format_table(table.get("table_data", "")),
                        "caption": table.get("caption", ""),
                        "source": table.get("source", ""),
                        "page": table.get("page", ""),
                        "description": table.get("description", "")
                    })
            
            # 13. Retourner les résultats finaux
            yield history, source_info, display_images(), display_tables()
        else:
            # Version sans streaming (code existant)
            result = rag_bot.chat(message, stream=False)
            history = history + [(message, result["response"])]
            
            # Préparer les informations sur les sources
            source_info = ""
            if "texts" in result:
                source_info += f"📚 {len(result['texts'])} textes • "
            if "images" in result:
                source_info += f"🖼️ {len(result['images'])} images • "
            if "tables" in result:
                source_info += f"📊 {len(result['tables'])} tableaux"
            
            if source_info:
                source_info = "Sources trouvées: " + source_info
                
            # Traiter les images et tableaux
            if show_sources and "images" in result and result["images"]:
                images = result["images"][:max_images]
                for img in images:
                    img_data = img.get("image_data")
                    if img_data:
                        image = base64_to_image(img_data)
                        if image:
                            current_images.append({
                                "image": image,
                                "caption": img.get("caption", ""),
                                "source": img.get("source", ""),
                                "page": img.get("page", ""),
                                "description": img.get("description", "")
                            })
            
            if show_sources and "tables" in result and result["tables"]:
                tables = result["tables"]
                for table in tables:
                    current_tables.append({
                        "data": rag_bot.format_table(table.get("table_data", "")),
                        "caption": table.get("caption", ""),
                        "source": table.get("source", ""),
                        "page": table.get("page", ""),
                        "description": table.get("description", "")
                    })
            
            return history, source_info, display_images(), display_tables()
    
    except Exception as e:
        error_msg = f"Une erreur est survenue: {str(e)}"
        traceback_text = traceback.format_exc()
        print(error_msg)
        print(traceback_text)
        history = history + [(message, error_msg)]
        return history, "Erreur lors du traitement de la requête", None, None

# Fonctions pour afficher les images et tableaux
def display_images():
    if not current_images:
        return None
    
    gallery = []
    for img_data in current_images:
        image = img_data["image"]
        if image:
            caption = f"{img_data['caption']} (Source: {img_data['source']}, Page: {img_data['page']})"
            gallery.append((image, caption))
    
    return gallery if gallery else None

def display_tables():
    if not current_tables:
        return None
    
    html = ""
    for table in current_tables:
        html += f"""
        <div style="margin-bottom: 20px; border: 1px solid #ddd; padding: 15px; border-radius: 8px;">
            <h3>{table['caption']}</h3>
            <p style="color:#666; font-size:0.9em;">Source: {table['source']}, Page: {table['page']}</p>
            <p><strong>Description:</strong> {table['description']}</p>
            <div style="background-color:#f5f5f5; padding:10px; border-radius:5px; overflow:auto;">
                <pre>{table['data']}</pre>
            </div>
        </div>
        """
    
    return html if html else None

# Fonction pour réinitialiser l'historique
def reset_conversation():
    global current_images, current_tables
    current_images = []
    current_tables = []
    
    rag_bot.clear_history()
    
    return [], "", None, None

# Interface Gradio
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
    gr.Markdown("# 📚 Assistant documentaire intelligent")
    
    with gr.Row():
        with gr.Column(scale=2):
            chat_interface = gr.Chatbot(
                height=600,
                show_label=False,
                layout="bubble",
                elem_id="chatbot"
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    show_label=False,
                    placeholder="Posez votre question...",
                    container=False,
                    scale=4
                )
                submit_btn = gr.Button("Envoyer", variant="primary", scale=1)
            
            clear_btn = gr.Button("Effacer la conversation")
            source_info = gr.Markdown("", elem_id="sources_info")
            
        with gr.Column(scale=1):
            with gr.Accordion("Options", open=True):
                # Sélecteur de modèle
                model_selector = gr.Dropdown(
                    choices=AVAILABLE_MODELS,
                    value=default_model,
                    label="Modèle Ollama",
                    info="Choisir le modèle de language à utiliser"
                )
                model_status = gr.Markdown(f"Modèle actuel: **{default_model}**")
                
                streaming = gr.Checkbox(
                    label="Mode streaming", 
                    value=True,
                    info="Voir les réponses s'afficher progressivement"
                )
                show_sources = gr.Checkbox(label="Afficher les sources", value=True)
                max_images = gr.Slider(
                    minimum=1,
                    maximum=10,
                    value=3,
                    step=1,
                    label="Nombre max d'images"
                )
            
            gr.Markdown("---")
            
            gr.Markdown("### 🖼️ Images pertinentes")
            image_gallery = gr.Gallery(
                label="Images pertinentes", 
                show_label=False,
                columns=2,
                height=300,
                object_fit="contain"
            )
            
            gr.Markdown("### 📊 Tableaux")
            tables_display = gr.HTML()
    
    # Connecter le changement de modèle
    model_selector.change(
        fn=change_model,
        inputs=model_selector,
        outputs=model_status
    )
    
    # Configuration des actions
    msg.submit(
        process_query,
        inputs=[msg, chat_interface, streaming, show_sources, max_images],
        outputs=[chat_interface, source_info, image_gallery, tables_display]
    ).then(lambda: "", outputs=msg)
    
    submit_btn.click(
        process_query,
        inputs=[msg, chat_interface, streaming, show_sources, max_images],
        outputs=[chat_interface, source_info, image_gallery, tables_display]
    ).then(lambda: "", outputs=msg)
    
    clear_btn.click(
        reset_conversation,
        outputs=[chat_interface, source_info, image_gallery, tables_display]
    )

    # Support amélioré pour les équations mathématiques avec KaTeX
    gr.Markdown("""
    <style>
    .gradio-container {max-width: 1200px !important}
    #chatbot {height: 600px; overflow-y: auto;}
    #sources_info {margin-top: 10px; color: #666;}
    
    /* Style pour les équations */
    .katex { font-size: 1.1em !important; }
    .math-inline { background: #f8f9fa; padding: 2px 5px; border-radius: 4px; }
    .math-display { background: #f8f9fa; margin: 10px 0; padding: 10px; border-radius: 5px; overflow-x: auto; text-align: center; }
    </style>
    
    <!-- Chargement de KaTeX -->
    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.8/dist/katex.min.css">
    <script src="https://cdn.jsdelivr.net/npm/katex@0.16.8/dist/katex.min.js"></script>
    <script src="https://cdn.jsdelivr.net/npm/katex@0.16.8/dist/contrib/auto-render.min.js"></script>
    
    <script>
    // Fonction pour rendre les équations avec KaTeX
    function renderMathInElement(element) {
        if (!window.renderMathInElement) return;
        
        window.renderMathInElement(element, {
            delimiters: [
                {left: '$$', right: '$$', display: true},
                {left: '$', right: '$', display: false},
                {left: '\\(', right: '\\)', display: false},
                {left: '\\[', right: '\\]', display: true}
            ],
            throwOnError: false,
            trust: true,
            strict: false
        });
    }
    
    // Fonction pour remplacer les underscores échappés qui posent problème
    function fixUnderscores(element) {
        const messages = element.querySelectorAll('.message');
        messages.forEach(msg => {
            const text = msg.innerHTML;
            // Remplacer les patterns comme u_(i) par u_{i} pour une meilleure compatibilité LaTeX
            const fixed = text.replace(/([a-zA-Z])_\(([^)]+)\)/g, '$1_{$2}');
            
            // Remplacer également les & qui peuvent causer des problèmes
            const cleanAmpersand = fixed.replace(/&amp;/g, '');
            
            if (text !== cleanAmpersand) {
                msg.innerHTML = cleanAmpersand;
            }
        });
    }
    
    // Observer les changements dans le chat
    function setupMathObserver() {
        const chatElement = document.getElementById('chatbot');
        if (!chatElement) {
            setTimeout(setupMathObserver, 500);
            return;
        }
        
        const observer = new MutationObserver((mutations) => {
            mutations.forEach(mutation => {
                if (mutation.type === 'childList' || mutation.type === 'subtree') {
                    const messages = chatElement.querySelectorAll('.message');
                    if (messages.length > 0) {
                        // D'abord corriger les underscores problématiques
                        fixUnderscores(chatElement);
                        
                        // Puis rendre les équations
                        messages.forEach(msg => {
                            renderMathInElement(msg);
                        });
                    }
                }
            });
        });
        
        observer.observe(chatElement, { 
            childList: true, 
            subtree: true,
            characterData: true
        });
        
        // Rendre les équations déjà présentes
        renderMathInElement(document);
    }
    
    // Initialisation lorsque la page est chargée
    document.addEventListener('DOMContentLoaded', function() {
        // Attendre que KaTeX soit chargé
        if (window.renderMathInElement) {
            setupMathObserver();
        } else {
            // Attendre le chargement de KaTeX
            document.querySelector('script[src*="auto-render.min.js"]').onload = setupMathObserver;
        }
    });
    </script>
    """)

if __name__ == "__main__":
    demo.queue()
    demo.launch(share=False, inbrowser=True)