import gradio as gr
import aiohttp
import asyncio
import random
import io
from PIL import Image
import os
from concurrent.futures import ThreadPoolExecutor


markdown_link = "For more generations, visit [Imagine Art](https://www.imagine.art/)."

TOKEN = os.getenv("TOKEN")
URL = os.getenv("URL")

if not URL:
    raise ValueError("URL environment variable is not set.")

async def async_generation_request(request_data, is_user_premium):
    seed = request_data.get("seed", random.randint(1, 10000000))

    headers = {
        "bearer": TOKEN,
        "Style-Id": request_data["model_id"]
    }
    if is_user_premium:
        headers["Premium"] = "True"

    for key, value in headers.items():
        if value is None:
            raise ValueError(f"Header value is missing for key: {key}")

    async with aiohttp.ClientSession() as session:
        for i in range(3):
            data = aiohttp.FormData()
            data.add_field('prompt', request_data["prompt"], content_type='multipart/form-data')
            data.add_field('style_id', request_data["model_id"])
            data.add_field('seed', str(seed))
            data.add_field('aspect_ratio', request_data["aspect_ratio"])
            data.add_field('negative_prompt', request_data["negative_prompt"])
            data.add_field('cfg', str(request_data["cfg_scale"]))

            async with session.post(URL, data=data, headers=headers) as response:
                if response.status == 200:
                    return await response.read()
                print(f"Failed API Call: {response.status} >>> {await response.text()}")

    return None

def run_async_in_thread(func, *args):
    with ThreadPoolExecutor() as executor:
        return executor.submit(asyncio.run, func(*args)).result()

def gradio_generation_request(prompt, aspect_ratio, negative_prompt, cfg_scale):
    request_data = {
        "prompt": prompt,
        "model_id": "32",
        "aspect_ratio": aspect_ratio,
        "negative_prompt": negative_prompt,
        "cfg_scale": cfg_scale,
        "seed": random.randint(1, 10000000)
    }

    output = run_async_in_thread(async_generation_request, request_data, True)  

    if output:
        return Image.open(io.BytesIO(output))
    else:
        return None

interface = gr.Interface(
    fn=gradio_generation_request,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Textbox(label="Aspect Ratio"),
        gr.Textbox(label="Negative Prompt"),
        gr.Slider(minimum=1, maximum=15, value=7.5, label="CFG Scale") 
    ],
    outputs="image",
    title='IMAGINE V4.1 GENERATION',
    description=markdown_link
)

interface.launch()