main.py

import discord
import os
import json
import io
import traceback
from dotenv import load_dotenv
from huggingface_hub import InferenceClient

SETTINGS_FILE = 'guild_settings.json'
IMAGE_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
TEXT_MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"


def load_settings():
    try:
        with open(SETTINGS_FILE, 'r') as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return {}


def save_settings(data):
    with open(SETTINGS_FILE, 'w') as f:
        json.dump(data, f, indent=4)


load_dotenv()
TOKEN = os.getenv('DISCORD_TOKEN')
HUGGINGFACE_TOKEN = os.getenv('HUGGINGFACE_TOKEN')

if not TOKEN or not HUGGINGFACE_TOKEN:
    raise ValueError("DISCORD_TOKEN and HUGGINGFACE_TOKEN must be set in the .env file.")

hf_client = InferenceClient(token=HUGGINGFACE_TOKEN)

intents = discord.Intents.default()
intents.messages = True
intents.message_content = True
intents.dm_messages = True

bot = discord.Bot(intents=intents)
guild_chat_channels = load_settings()


async def get_ai_response(prompt: str, author_name: str):
    if not prompt.strip():
        return None

    print(f"User '{author_name}' is prompting text generator with: '{prompt}'")

    messages = [
        {"role": "system",
         "content": "You are a friendly and helpful chatbot named EchoAI. Your job is to be a friend to everyone who talks to you and to help them out with whatever they need."},
        {"role": "user", "content": prompt},
    ]

    try:
        response = hf_client.chat_completion(
            messages=messages,
            model=TEXT_MODEL_ID,
            max_tokens=250,
            temperature=0.8,
            stream=False,
        )
        return response.choices[0].message.content

    except Exception as e:
        print("An error occurred during text generation:")
        traceback.print_exc()
        return "Sorry, I'm having trouble connecting to the AI service right now. It might be busy or loading. Please try again in a moment."


@bot.event
async def on_ready():
    print(f'Logged in as {bot.user}')
    print(f'Bot is completely operational!')
    print('------')


@bot.event
async def on_message(message: discord.Message):
    if message.author == bot.user:
        return

    if isinstance(message.channel, discord.DMChannel):
        async with message.channel.typing():
            response = await get_ai_response(message.content, message.author.name)
            if response: await message.channel.send(response)
        return

    if message.guild:
        was_mentioned = bot.user in message.mentions

        guild_id_str = str(message.guild.id)
        chat_channel_id = guild_chat_channels.get(guild_id_str)
        is_in_chat_channel = chat_channel_id and message.channel.id == int(chat_channel_id)

        if was_mentioned or is_in_chat_channel:
            prompt = message.content.replace(f'<@{bot.user.id}>', '').strip()

            async with message.channel.typing():
                response = await get_ai_response(prompt, message.author.name)
                if response:
                    if was_mentioned:
                        await message.reply(response)
                    else:
                        await message.channel.send(response)


@bot.slash_command(
    name="set_chat_channel",
    description="[ADMIN] Set this channel as the primary channel for conversation."
)
@discord.commands.default_permissions(manage_guild=True)
async def set_chat_channel(ctx: discord.ApplicationContext):
    guild_id = str(ctx.guild.id)
    channel_id = str(ctx.channel.id)
    guild_chat_channels[guild_id] = channel_id
    save_settings(guild_chat_channels)
    await ctx.respond(f"✅ **AI Chat Enabled!** I will now chat with users in {ctx.channel.mention}.", ephemeral=True)


@bot.slash_command(
    name="generate_image",
    description="Generate an image from a text prompt (works in DMs too).",
    dm_permission=True
)
async def generate_image(
        ctx: discord.ApplicationContext,
        prompt: discord.Option(str, "Describe the image you want to create.", required=True)
):
    await ctx.defer()

    try:
        image_pil = hf_client.text_to_image(prompt, model=IMAGE_MODEL_ID)
        with io.BytesIO() as image_binary:
            image_pil.save(image_binary, 'PNG')
            image_binary.seek(0)
            embed = discord.Embed(title="🎨 Image Generated 🎨", color=discord.Color.blue())
            embed.add_field(name="Your Prompt", value=prompt, inline=False)
            embed.set_footer(text=f"Requested by {ctx.author.name}")
            embed.set_image(url="attachment://image.png")
            await ctx.followup.send(embed=embed, file=discord.File(fp=image_binary, filename='image.png'))
    except Exception as e:
        print(f"Error generating image: {e}")
        traceback.print_exc()
        await ctx.followup.send("😭 Sorry, I couldn't generate the image. Please try again in a moment.")


if __name__ == "__main__":
    bot.run(TOKEN)