from huggingface_hub import InferenceClient
from typing import List, Tuple
import gradio as gr

# Ensure you have the required libraries installed

# Chatbot configuration
class ChatConfig:
    MODEL = "google/gemma-3-27b-it"
    # MODEL = "mistralai/Mistral-7B-Instruct-v0.1"
    # MODEL = "google/gemma-2-9b-it"
    DEFAULT_SYSTEM_MSG = "You are a extremely smart and useful Chatbot."
    DEFAULT_MAX_TOKENS = 1024
    DEFAULT_TEMP = 0.4
    DEFAULT_TOP_P = 0.95

import os
HF_TOKEN = os.getenv("HUGGING_FACE_HUB_TOKEN")

# Not recommended to hardcode your token in the script for security reasons.
# HF_TOKEN = "hf_YOUR_TOKEN_HERE"

client = InferenceClient(ChatConfig.MODEL, token=HF_TOKEN)

def generate_response(
    message: str,
    history: List[Tuple[str, str]],
    system_message: str = ChatConfig.DEFAULT_SYSTEM_MSG,
    max_tokens: int = ChatConfig.DEFAULT_MAX_TOKENS,
    temperature: float = ChatConfig.DEFAULT_TEMP,
    top_p: float = ChatConfig.DEFAULT_TOP_P
) -> str:
    
    messages = [{"role": "system", "content": system_message}]
    
    # Conversation history
    for user_msg, bot_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})
    
    messages.append({"role": "user", "content": message})
    
    response = ""

    # Stream the response
    # The stream=True parameter allows for real-time response generation
    for chunk in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = chunk.choices[0].delta.content or ""
        response += token
        yield response


def create_interface() -> gr.ChatInterface:
    """Create and configure the chat interface."""
    # Custom CSS for a modern look
    custom_css = """
    body {
        background: linear-gradient(to right, #a8c0ff, #fafdff);
        font-family: 'Comic Sans MS', sans-serif;
    }

    .chatbot .message { 
        border-radius: 12px; 
        margin: 5px; 
        padding: 10px; 
    }
    .chatbot .user-message { 
        background-color: #b3e5fc;
        color: #0d47a1; 
    }
    .chatbot .bot-message { 
        background-color: #e3f2fd;
        color: #0d47a1;
    }
    .gr-button { 
        border-radius: 8px; 
        padding: 8px 16px; 
    }
    """

    with gr.Blocks(css=custom_css) as app:
        gr.Markdown("# Your Gemma Three-(end) 💎")

        # Custom chatbot 
        chatbot = gr.Chatbot(
            label="Gemma Chat",
            avatar_images=("./user.png", "./gemma.png"),
            height=450,
            show_copy_button=True
        )

        # Chat interface
        interface = gr.ChatInterface(
            fn=generate_response,
            chatbot=chatbot,
            title="Your Gemma 3-(iend)",
            theme=gr.themes.Soft(),
            css=custom_css,
            additional_inputs=[
                gr.Textbox(
                    value=ChatConfig.DEFAULT_SYSTEM_MSG,
                    label="System Prompt (You can change the text below)",
                    lines=2,
                    placeholder="Enter system message..."
                ),
                gr.Slider(
                    minimum=1,
                    maximum=8192,
                    value=ChatConfig.DEFAULT_MAX_TOKENS,
                    step=1,
                    label="Max Tokens (Response Length)",
                    info="Controls response length"
                ),
                gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=ChatConfig.DEFAULT_TEMP,
                    step=0.1,
                    label="Temperature (0.1-1.0)",
                    info="Controls randomness"
                ),
                gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=ChatConfig.DEFAULT_TOP_P,
                    step=0.05,
                    label="Top-P (0.1-1.0)",
                    info="Controls diversity"
                )
            ],

            additional_inputs_accordion=gr.Accordion(label="Advanced Settings", open=False)
        )
    
    return app


# def create_interface() -> gr.Blocks:
#     """Create and configure the chat interface."""
#     # Custom CSS for a modern look
#     custom_css = """
#     body {
#         background: linear-gradient(to right, #a8c0ff, #fafdff);
#         font-family: 'Comic Sans MS', sans-serif;
#     }

#     .chatbot .message { 
#         border-radius: 12px; 
#         margin: 5px; 
#         padding: 10px; 
#     }
#     .chatbot .user-message { 
#         background-color: #b3e5fc;
#         color: #0d47a1; 
#     }
#     .chatbot .bot-message { 
#         background-color: #e3f2fd;
#         color: #0d47a1;
#     }
#     .gr-button { 
#         border-radius: 8px; 
#         padding: 8px 16px; 
#     }
#     """
    
#     # JavaScript pour jouer un son à la soumission
#     custom_js = """
#     function setupSubmitSound() {
#         // Create the audio element for the sound
#         const submitSound = new Audio('/gradio_api/file=submit_sound.mp3');
        
#         // Fonction pour attacher l'écouteur d'événement
#         function attachSoundToButton() {
#             // Rechercher le bouton de soumission dans l'interface
#             const submitButton = document.querySelector('button.send-btn') || 
#                                  document.querySelector('button[aria-label="Submit"]') ||
#                                  document.querySelector('.submit');
            
#             if (submitButton) {
#                 console.log('Found submit button, attaching sound');
#                 submitButton.addEventListener('click', function() {
#                     submitSound.currentTime = 0; // Réinitialiser le son
#                     submitSound.play();
#                 });
#             } else {
#                 // Si le bouton n'est pas encore disponible, réessayer plus tard
#                 console.log('Button not found, trying again in 500ms');
#                 setTimeout(attachSoundToButton, 500);
#             }
#         }
        
#         // Démarrer la recherche du bouton
#         attachSoundToButton();
#     }
    
#     // Appeler la fonction d'initialisation
#     setupSubmitSound();
#     """

#     with gr.Blocks(css=custom_css, js=custom_js) as app:
#         gr.Markdown("# Your Gemma Three-(end) 💎")

#         # Custom chatbot 
#         chatbot = gr.Chatbot(
#             label="Gemma Chat",
#             avatar_images=("./user.png", "./gemma.png"),
#             height=450,
#             show_copy_button=True
#         )

#         # Chat interface
#         interface = gr.ChatInterface(
#             fn=generate_response,
#             chatbot=chatbot,
#             title="Your Gemma 3-(iend)",
#             theme=gr.themes.Soft(),
#             additional_inputs=[
#                 gr.Textbox(
#                     value=ChatConfig.DEFAULT_SYSTEM_MSG,
#                     label="System Prompt (You can change the text below)",
#                     lines=2,
#                     placeholder="Enter system message..."
#                 ),
#                 gr.Slider(
#                     minimum=1,
#                     maximum=8192,
#                     value=ChatConfig.DEFAULT_MAX_TOKENS,
#                     step=1,
#                     label="Max Tokens (Response Length)",
#                     info="Controls response length"
#                 ),
#                 gr.Slider(
#                     minimum=0.1,
#                     maximum=1.0,
#                     value=ChatConfig.DEFAULT_TEMP,
#                     step=0.1,
#                     label="Temperature (0.1-1.0)",
#                     info="Controls randomness"
#                 ),
#                 gr.Slider(
#                     minimum=0.1,
#                     maximum=1.0,
#                     value=ChatConfig.DEFAULT_TOP_P,
#                     step=0.05,
#                     label="Top-P (0.1-1.0)",
#                     info="Controls diversity"
#                 )
#             ],
#             additional_inputs_accordion=gr.Accordion(label="Advanced Settings", open=False)
#         )
    
#     return app

def main():
    app = create_interface()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_api=False,
        show_error=True,
        debug=True
    )

if __name__ == "__main__":
    main()