# bot.py
import discord
from gradio_client import Client
from huggingface_hub import InferenceClient
import os

# Get tokens from environment variables (Hugging Face Spaces secrets)
TOKEN = os.getenv("DISCORD_TOKEN")
HF_TOKEN = os.getenv("HF_TOKEN")

# Check if tokens are available
if not TOKEN or not HF_TOKEN:
    raise ValueError("DISCORD_TOKEN and HF_TOKEN must be set as environment variables in Spaces secrets")

# Configuration des intents Discord
intents = discord.Intents.default()
intents.message_content = True

# Initialisation des clients
discord_client = discord.Client(intents=intents)
hf_client = InferenceClient(api_key=HF_TOKEN)

# Fonction pour obtenir une réponse du modèle
async def get_model_response(message):
    messages = [
        {"role": "user", "content": message}
    ]
    
    response = ""
    stream = hf_client.chat.completions.create(
        model="Qwen/Qwen2.5-72B-Instruct",
        messages=messages,
        temperature=0.5,
        max_tokens=2048,
        top_p=0.7,
        stream=True
    )
    
    for chunk in stream:
        if chunk.choices[0].delta.content:
            response += chunk.choices[0].delta.content
    
    return response

@discord_client.event
async def on_ready():
    print(f'Connecté en tant que {discord_client.user}')

@discord_client.event
async def on_message(message):
    # Ignorer les messages du bot lui-même
    if message.author == discord_client.user:
        return

    # Vérifier si le bot est mentionné
    if discord_client.user in message.mentions:
        try:
            # Extraire le contenu du message sans la mention
            clean_message = message.content.replace(f"<@{discord_client.user.id}>", "").strip()
            
            if not clean_message:
                await message.channel.send("Veuillez inclure un message après ma mention !")
                return

            # Obtenir la réponse du modèle
            async with message.channel.typing():
                response = await get_model_response(clean_message)
                
                # Diviser la réponse si elle dépasse la limite Discord (2000 caractères)
                if len(response) > 2000:
                    parts = [response[i:i+2000] for i in range(0, len(response), 2000)]
                    for part in parts:
                        await message.channel.send(part)
                else:
                    await message.channel.send(response)
                    
        except Exception as e:
            await message.channel.send(f"Une erreur est survenue : {str(e)}")

# Lancer le bot
if __name__ == "__main__":
    # Vérifier que les tokens sont définis
    if TOKEN == "VOTRE_TOKEN_DISCORD" or HF_TOKEN == "VOTRE_TOKEN_HF":
        print("Veuillez remplacer les tokens par vos propres tokens Discord et Hugging Face")
    else:
        discord_client.run(TOKEN)