import gradio as gr
import requests
import json

def respond(message, history, system_message, max_tokens, temperature, top_p):
    """
    Fonction pour générer une réponse en utilisant Ollama
    
    Args:
    - message: Le dernier message de l'utilisateur
    - history: Historique des conversations
    - system_message: Message système pour guider le comportement du modèle
    - max_tokens: Nombre maximal de tokens à générer
    - temperature: Contrôle la créativité de la réponse
    - top_p: Échantillonnage nucléaire
    
    Returns:
    - Réponse générée par le modèle
    """
    # Préparer les messages pour le contexte
    messages = [{"role": "system", "content": system_message}]
    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})
    
    # Préparer les données pour la requête Ollama
    data = {
        "model": "hf.co/ibrahimBlyc/LA_Llama:latest",  # Assurez-vous que ce modèle est disponible localement
        "messages": messages,
        "stream": True,
        "options": {
            "num_predict": max_tokens,
            "temperature": temperature,
            "top_p": top_p
        }
    }
    
    # URL du serveur Ollama local
    url = "http://localhost:11434/api/chat"
    
    try:
        # Envoyer la requête en streaming
        response = ""
        with requests.post(url, json=data, stream=True) as stream_response:
            stream_response.raise_for_status()
            for chunk in stream_response.iter_lines():
                if chunk:
                    try:
                        # Décoder correctement le chunk JSON
                        chunk_decoded = chunk.decode('utf-8')
                        chunk_json = json.loads(chunk_decoded)
                        
                        # Extraire le token de la réponse
                        if 'message' in chunk_json and 'content' in chunk_json['message']:
                            token = chunk_json['message']['content']
                            response += token
                            yield response
                        
                    except json.JSONDecodeError as e:
                        print(f"Erreur de décodage JSON : {e}")
                        yield f"Erreur : Impossible de décoder la réponse du serveur."
                        return
                    except Exception as e:
                        print(f"Erreur inattendue : {e}")
                        yield f"Erreur : Une erreur inattendue s'est produite."
                        return
    
    except requests.exceptions.RequestException as e:
        print(f"Erreur de requête : {e}")
        yield f"Erreur : Impossible de communiquer avec le serveur Ollama."

# Créer l'interface Gradio
demo = gr.ChatInterface(
    respond,
    title="Ollama Chatbot",
    description="Discutez avec un modèle Ollama local",
    additional_inputs=[
        gr.Textbox(
            value="Tu es un assistant IA amical et utile qui répond de manière concise et informative.", 
            label="Message système",
            lines=3
        ),
        gr.Slider(
            minimum=1, 
            maximum=2048, 
            value=512, 
            step=1, 
            label="Nombre max de tokens"
        ),
        gr.Slider(
            minimum=0.1, 
            maximum=2.0, 
            value=0.7, 
            step=0.1, 
            label="Température"
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (échantillonnage nucléaire)"
        ),
    ],
    theme="default"
)

# Lancement de l'application
if __name__ == "__main__":
    # Vérifier que le serveur Ollama est bien installé
    try:
        import ollama
        print("Serveur Ollama détecté. Lancement de l'interface...")
        demo.launch(share=True)
    except ImportError:
        print("Le serveur Ollama n'est pas installé. Veuillez l'installer.")
    except Exception as e:
        print(f"Erreur lors du lancement : {e}")