import asyncio
import os
import threading
import random
from threading import Event
from typing import Optional

import discord
import gradio as gr
from discord import Permissions
from discord.ext import commands
from discord.utils import oauth_url

import gradio_client as grc
from gradio_client.utils import QueueError
#---------------------------------------------------------------------------------------------------------------------
event = Event()

DISCORD_TOKEN = os.getenv("DISCORD_TOKEN")
#---------------------------------------------------------------------------------------------------------------------
async def wait(job):
    while not job.done():
        await asyncio.sleep(0.2)
#---------------------------------------------------------------------------------------------------------------------
def get_client(session: Optional[str] = None) -> grc.Client:
    client = grc.Client("huggingface-projects/audioldm2-text2audio-text2music", hf_token=os.getenv("HF_TOKEN"))
    if session:
        client.session_hash = session
    return client
#---------------------------------------------------------------------------------------------------------------------
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} (ID: {bot.user.id})")
    synced = await bot.tree.sync()
    print(f"Synced commands: {', '.join([s.name for s in synced])}.")
    event.set()
    print("------")
#---------------------------------------------------------------------------------------------------------------------
@bot.hybrid_command(
    name="audioldm2",
    description="Enter a prompt to generate music!",
)
async def audioldm2(ctx, prompt: str):
    """Audioldm2 generation"""
    try:
        await audioldm2_create(ctx, prompt)
    except Exception as e:
        print(f"Error: (app.py){e}")
#---------------------------------------------------------------------------------------------------------------------
async def audioldm2_create(ctx, prompt):
    """Runs audioldm2_create_job in executor"""
    try:
        message = await ctx.send(f"**{prompt}** - {ctx.author.mention} Generating...")
        thread = await message.create_thread(name=prompt[:100])

        loop = asyncio.get_running_loop()
        client = await loop.run_in_executor(None, get_client, None)
        job = client.submit(prompt, fn_index=0)
        await wait(job)
        
        try:
            job.result()
            video = job.outputs()[0]
        except QueueError:
            await thread.send("The gradio space powering this bot is really busy! Please try again later!")

        short_filename = prompt[:20]
        video_filename = f"{short_filename}.mp4"

        with open(video, "rb") as file:
            discord_video_file = discord.File(file, filename=video_filename)
        await thread.send(file=discord_video_file)

    except Exception as e:
        print(f"audioldm2_create Error: {e}")
#---------------------------------------------------------------------------------------------------------------------
def run_bot():
    if not DISCORD_TOKEN:
        print("DISCORD_TOKEN NOT SET")
        event.set()
    else:
        bot.run(DISCORD_TOKEN)


threading.Thread(target=run_bot).start()

event.wait()

with gr.Blocks() as demo:
    gr.Markdown(
    """
    # Discord bot of https://huggingface.co/spaces/facebook/MusicGen
    https://discord.com/api/oauth2/authorize?client_id=1151888750662664222&permissions=309237696512&scope=bot
    """
    )

demo.launch()