import discord
import os
import gradio as gr
import asyncio
import time # 
from discord.ext import commands
from gradio_client import Client
from PIL import Image

breaking bot

DFIF_TOKEN = os.getenv('HF_TOKEN')
DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)

jojogan = Client("akhaliq/JoJoGAN", DFIF_TOKEN)

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

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

MAX_WORKERS = 5

worker_semaphore = asyncio.Semaphore(MAX_WORKERS)

async def jojo(ctx):
    start_time = time.time()  
    style = 'JoJo'
    atchurl = 'https://cdn.discordapp.com/attachments/1100458786826747945/1111746037640601610/image.png'
    im = jojogan.predict(atchurl, style)
    end_time = time.time()
    generation_time = end_time - start_time    
    await ctx.send(f"jojo image generated in {generation_time:.2f} seconds.")
    await ctx.send(file=discord.File(im))

async def process_requests(queue):
    while True:
        ctx = await queue.get()
        async with worker_semaphore:
            await jojo(ctx)
        queue.task_done()

@bot.command()
async def command(ctx, num_requests: int):
    start_time = time.time() 
    queue = asyncio.Queue()

    # Start the worker tasks
    worker_tasks = []
    for _ in range(MAX_WORKERS):
        task = asyncio.create_task(process_requests(queue))
        worker_tasks.append(task)

    # Add the requests to the queue
    for _ in range(num_requests):
        await queue.put(ctx)

    # Wait for all requests to be processed
    await queue.join()

    # Cancel the worker tasks
    for task in worker_tasks:
        task.cancel()

    await ctx.send("Command executed.")
    end_time = time.time()
    generation_time = end_time - start_time    
    await ctx.send(f"{style} everything done in {generation_time:.2f} seconds.")    
    
def run_bot():
    bot.run(DISCORD_TOKEN)


async def run_gradio_interface():
    def greet(name):
        return "Hello " + name + "!"

    demo = gr.Interface(fn=greet, inputs="text", outputs="text")
    await demo.launch()


async def main():
    loop = asyncio.get_event_loop()
    tasks = [
        loop.create_task(run_bot()),
        loop.create_task(run_gradio_interface())
    ]
    await asyncio.wait(tasks)


asyncio.run(main())