import gradio as gr
from huggingface_hub import InferenceClient
import base64
from PIL import Image
import io

# Client d'inférence pour Apriel-1.6-15b-Thinker
model_id = "ServiceNow-AI/Apriel-1.6-15b-Thinker"
client = InferenceClient(model=model_id)

def encode_image_to_base64(image):
    """Encode l'image en base64 pour l'API"""
    if image is None:
        return None
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return f"data:image/png;base64,{img_str}"

def chat_with_apriel(message, history, image=None, temperature=0.6, max_tokens=1024):
    """Fonction de chatbot avec support texte et image"""
    if not message:
        return ""
    try:
        messages = []
        for user_msg, assistant_msg in history:
            messages.append({"role": "user", "content": user_msg})
            messages.append({"role": "assistant", "content": assistant_msg})
        if image is not None:
            img_b64 = encode_image_to_base64(image)
            messages.append({"role": "user", "content": [{"type": "text", "text": message}, {"type": "image_url", "image_url": {"url": img_b64}}]})
        else:
            messages.append({"role": "user", "content": message})
        response = client.chat_completion(messages=messages, max_tokens=max_tokens, temperature=temperature, stream=False)
        return response.choices[0].message.content
    except Exception as e:
        return f"⚠️ Erreur: {str(e)}\n\nLe modèle peut prendre ~30s au premier démarrage."

def generate_text(prompt, temperature=0.6, max_tokens=1024):
    """Fonction de génération de texte"""
    if not prompt:
        return ""
    try:
        messages = [{"role": "user", "content": prompt}]
        response = client.chat_completion(messages=messages, max_tokens=max_tokens, temperature=temperature, stream=False)
        return response.choices[0].message.content
    except Exception as e:
        return f"⚠️ Erreur: {str(e)}\n\nLe modèle peut prendre ~30s au premier démarrage."

with gr.Blocks(title="Apriel Chatbot Complete") as demo:
    gr.Markdown("""
    # 🤖 Apriel-1.6-15B-Thinker - Interface Complète
    Modèle multimodal 15B paramètres - ServiceNow AI
    🚀 **Utilise l'API Inference HuggingFace** - Rapide & performant!
    """)
    with gr.Tabs():
        with gr.Tab("💬 Chatbot"):
            gr.Markdown("### Chat interactif avec Apriel")
            with gr.Row():
                with gr.Column(scale=2):
                    chatbot = gr.Chatbot(label="Conversation", height=500)
                    with gr.Row():
                        msg = gr.Textbox(label="Message", placeholder="Votre message...", scale=4)
                    with gr.Row():
                        image_input = gr.Image(label="Image (optionnelle)", type="pil", height=200)
                    with gr.Row():
                        clear = gr.Button("🗑️ Effacer")
                        submit = gr.Button("📤 Envoyer", variant="primary")
                with gr.Column(scale=1):
                    gr.Markdown("### ⚙️ Paramètres")
                    temperature_chat = gr.Slider(0.1, 1.0, 0.6, 0.1, label="Température")
                    max_tokens_chat = gr.Slider(128, 2048, 1024, 128, label="Tokens max")
                    gr.Markdown("""
                    ### 📚 Capacités
                    - Raisonnement complexe
                    - Images & texte
                    - Code & math
                    - Conversations
                    
                    **Note:** 1ère réponse ~30s
                    """)
            def respond(message, chat_history, image, temp, max_tok):
                if not message:
                    return "", chat_history, None
                bot_message = chat_with_apriel(message, chat_history, image, temp, max_tok)
                chat_history.append((message, bot_message))
                return "", chat_history, None
            submit.click(respond, [msg, chatbot, image_input, temperature_chat, max_tokens_chat], [msg, chatbot, image_input])
            msg.submit(respond, [msg, chatbot, image_input, temperature_chat, max_tokens_chat], [msg, chatbot, image_input])
            clear.click(lambda: (None, None), None, [chatbot, image_input], queue=False)
        with gr.Tab("✨ Génération"):
            gr.Markdown("### Génération de texte")
            with gr.Row():
                with gr.Column(scale=2):
                    prompt_input = gr.Textbox(label="Prompt", placeholder="Entrez votre prompt...", lines=5)
                    generate_btn = gr.Button("🚀 Générer", variant="primary")
                    output_text = gr.Textbox(label="Réponse", lines=15)
                with gr.Column(scale=1):
                    gr.Markdown("### ⚙️ Paramètres")
                    temperature_gen = gr.Slider(0.1, 1.0, 0.6, 0.1, label="Température")
                    max_tokens_gen = gr.Slider(128, 2048, 1024, 128, label="Tokens max")
                    gr.Markdown("""
                    ### 💡 Exemples
                    - "Explique le concept de..."
                    - "Écris une fonction Python..."
                    - "Résous ce problème math..."
                    """)
            generate_btn.click(generate_text, inputs=[prompt_input, temperature_gen, max_tokens_gen], outputs=output_text)
        with gr.Tab("ℹ️ À propos"):
            gr.Markdown("""
            ## Apriel-1.6-15B-Thinker
            
            ### 🎯 Caractéristiques
            - **15B paramètres** - Compact mais puissant
            - **Multimodal** - Texte + Images
            - **Raisonnement avancé**
            - **30% plus efficace** que v1.5
            - **Score 57** sur Artificial Analysis
            
            ### 📊 Performance
            - MMMU: 72% | MathVista: 79.9%
            - AIME 25: 88% | IFBench: 69%
            
            ### 🚀 Architecture
            Utilise **l'API Inference HuggingFace**:
            - 💨 Réponses rapides
            - ⚙️ Infrastructure GPU optimisée
            - 🌐 Disponibilité maximale
            
            ### 🔗 Liens
            - [Modèle HF](https://huggingface.co/ServiceNow-AI/Apriel-1.6-15b-Thinker)
            - [Blog Post](https://huggingface.co/blog/ServiceNow-AI/apriel-1p6-15b-thinker)
            
            MIT License - ServiceNow AI
            """)

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())