import os
import discord
from discord.ext import commands
from huggingface_hub import InferenceClient

# --- Hugging Face Client ---
client = InferenceClient(
    provider="cerebras",
    api_key=os.getenv("HF_TOKEN"),
)

# --- Discord Bot Setup ---
intents = discord.Intents.default()
intents.message_content = True

bot = commands.Bot(command_prefix="!", intents=intents)

@bot.event
async def on_ready():
    print(f"✅ Logged in as {bot.user}")

# --- Chat Command ---
@bot.command(name="ask")
async def ask(ctx, *, prompt: str):
    try:
        completion = client.chat.completions.create(
            model="meta-llama/Llama-3.3-70B-Instruct",
            messages=[
                {"role": "user", "content": prompt}
            ],
            max_tokens=500,
        )

        reply = completion.choices[0].message.content
        await ctx.reply(reply)

    except Exception as e:
        await ctx.reply(f"⚠️ Error: {e}")

# --- Run Bot ---
bot.run(os.getenv("DISCORD_TOKEN"))