import asyncio
import io
import json
import os

import aiohttp
from PIL import Image
from dotenv import load_dotenv
load_dotenv()
import PIL

API_URL = os.getenv("API_URL")
API_KEY = os.getenv("API_KEY")

import gradio as gr


async def load_image_from_url(url):
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            image_data = await response.read()
            return Image.open(io.BytesIO(image_data))

async def greet(prompt):

    url = API_URL
    headers = {
        'x-qrbtf-key': f'{API_KEY}',
    }

    full_response: str = ""

    payload = {
        'url': 'https://qrbtf.com/',
        'prompt': prompt,
    }

    async with aiohttp.ClientSession(headers=headers) as session:

        async with session.post(url, json=payload, ssl=False) as response:

            await asyncio.sleep(0)

            async for (chunk, _) in response.content.iter_chunks():
                chunk = json.loads(chunk.decode("utf-8"))
                if chunk["type"] == "result":
                    data = chunk["data"]
                    url = data["download_url"]
                    print(url)
                    return await load_image_from_url(url)

    return full_response


with gr.Blocks() as demo:
    gr.Markdown("""
    
    # QRBTF.AI API Demo
    
    - [Join Discord](https://discord.gg/V9CNuqYfte)
    - [Official website](https://qrbtf.com/)
    
    """)
    with gr.Row():
        with gr.Column():
            url = gr.Textbox(
                label="URL",
                placeholder="https://",
                value="https://qrbtf.com/",
                interactive=False
            )
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Enter a prompt here",
            )
            with gr.Row():
                clear_btn = gr.Button(
                    "Clear",
                )
                btn = gr.Button(
                    "Call API",
                    variant="primary"
                )
        with gr.Column():
            out = gr.Image()
    btn.click(fn=greet, inputs=prompt, outputs=out)
    clear_btn.click(fn=lambda: prompt.update(""))

demo.launch()