import asyncio
import glob
import os
import random
import threading
from discord.ext import commands
import discord
import gradio as gr
from gradio_client import Client


HF_TOKEN = os.getenv("HF_TOKEN")
wuerstchen_client = Client("huggingface-projects/Wuerstchen-duplicate", HF_TOKEN)
DISCORD_TOKEN = os.getenv("DISCORD_TOKEN")


intents = discord.Intents.all()
bot = commands.Bot(command_prefix="/", intents=intents)


@bot.event
async def on_ready():
    print(f"Logged in as {bot.user} (ID: {bot.user.id})")
    synced = await bot.tree.sync()
    print(f"Synced commands: {', '.join([s.name for s in synced])}.")
    print("------")


@bot.hybrid_command(
    name="wuerstchen",
    description="Enter a prompt to generate art!",
)
async def wuerstchen_command(ctx, prompt: str):
    """Wuerstchen generation"""
    try:
        await run_wuerstchen(ctx, prompt)
    except Exception as e:
        print(f"Error wuerstchen: (app.py){e}")


def wuerstchen_inference(prompt):
    """Inference for Wuerstchen"""
    negative_prompt = ""
    seed = random.randint(0, 1000)
    width = 1024
    height = 1024
    prior_num_inference_steps = 60
    prior_guidance_scale = 4
    decoder_num_inference_steps = 12
    decoder_guidance_scale = 0
    num_images_per_prompt = 1

    result_path = wuerstchen_client.predict(
        prompt,
        negative_prompt,
        seed,
        width,
        height,
        prior_num_inference_steps,
        prior_guidance_scale,
        decoder_num_inference_steps,
        decoder_guidance_scale,
        num_images_per_prompt,
        api_name="/run",
    )
    png_file = list(glob.glob(f"{result_path}/**/*.png"))
    return png_file[0]


async def run_wuerstchen(ctx, prompt):
    """Responds to /Wuerstchen command"""
    try:
        if ctx.guild.id == 879548962464493619:
            if ctx.channel.id != 1151792944676864041:
                return
        message = await ctx.send(f"**{prompt}** - {ctx.author.mention} (generating...)")

        loop = asyncio.get_running_loop()
        result_path = await loop.run_in_executor(None, wuerstchen_inference, prompt)

        await message.delete()
        with open(result_path, "rb") as f:
            await ctx.channel.send(f"**{prompt}** - {ctx.author.mention}", file=discord.File(f, "wuerstchen.png"))
    except Exception as e:
        print(f"Error: {e}")


def run_bot():
    bot.run(DISCORD_TOKEN)


threading.Thread(target=run_bot).start()
"""This allows us to run the Discord bot in a Python thread"""


welcome_message = """
## Add this bot to your server by clicking this link: 

https://discord.com/api/oauth2/authorize?client_id=1155489509518098565&permissions=51200&scope=bot

## How to use it?

The bot can be triggered via `/wuerstchen` followed by your text prompt.

This will generate an image based on your prompt, which is then posted in the channel!

⚠️ Note ⚠️: Please make sure this bot's command does have the same name as another command in your server.

⚠️ Note ⚠️: Bot commands do not work in DMs with the bot as of now.
"""


with gr.Blocks() as demo:
    gr.Markdown(f"""
    # Discord bot of https://huggingface.co/spaces/warp-ai/Wuerstchen
    {welcome_message}
    """)


demo.queue(concurrency_count=100)
demo.queue(max_size=100)
demo.launch()