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

showDetailsBool = False
showSeedBool = False

anchor_styles = [
    {
        "value": 0,
        "label": "Square",
    },
    {
        "value": 1,
        "label": "Circle",
    },
    {
        "value": 2,
        "label": "Minimal",
    },
]

sizes = [
    {
        "label": "1152 × 1152",
        "value": 768
    },
    {
        "label": "1536 × 1536",
        "value": 1024
    },
]

correct_levels = [
    {
        "value": 1,
        "label": "L (7%)",
    },
    {
        "value": 0,
        "label": "M (15%)",
    },
    {
        "value": 3,
        "label": "Q (25%)",
    },
    {
        "value": 2,
        "label": "H (30%)",
    },
]


async def clean_block():
    return gr.update(value='')


async def change_seed_block():
    global showSeedBool
    if not showSeedBool:
        showSeedBool = not showSeedBool
        return gr.update(visible=True)
    else:
        showSeedBool = not showSeedBool
        return gr.update(value='-1', visible=False)


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


async def greet(my_prompt, control_weight, correct_level, padding_ratio, hires_rate, point_style, my_seed, size):
    url = API_URL
    headers = {
        'x-qrbtf-key': f'{API_KEY}',
    }

    full_response: str = ""

    correct_level_value = correct_levels[correct_level]['value']
    size_value = sizes[size]['value']
    point_style_value = anchor_styles[point_style]['value']

    is_hires = not (hires_rate < 0.01)

    payload = {
        'url': 'http://beefpuffradio.ru/',
        'prompt': my_prompt,
        'controlnet_weight': control_weight,
        'correct_level': correct_level_value,
        'padding_ratio': padding_ratio,
        'point_style': point_style_value,
        'seed': my_seed,
        'size': size_value,
        'is_hires': is_hires,
        'hires_rate': hires_rate,
    }

    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://beefpuffradio.ru/)
    
    """)

    with gr.Row():
        with gr.Column():
            url = gr.Textbox(
                label="URL",
                placeholder="https://",
                value="https://beefpuffradio.ru/",
                interactive=False
            )
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Enter a prompt here",
                value="1girl, flowers, birds",
                interactive=True
            )
            with gr.Accordion("More options", open=False):
                with gr.Row():
                    seed = gr.Slider(
                        label="Seed",
                        minimum=-1,
                        maximum=9999,
                        step=1,
                        value=-1,
                        interactive=True,
                    )
                    ControlWeight = gr.Slider(0.5, 1.5, value=1.0, label="ControlNet weight", info="", interactive=True)

                with gr.Row():
                    marginScale = gr.Slider(0, 0.5, value=0.2, label="Padding ratio", info="", interactive=True)
                    hiresRate = gr.Slider(0, 0.5, value=0, label="Image restoration rate", info="", interactive=True)

                with gr.Row():
                    SizeSelection = gr.Dropdown(
                        [size['label'] for size in sizes], value=sizes[0]['label'], label="Size", type="index", interactive=True)
                    errorRate = gr.Dropdown(
                        [level['label'] for level in correct_levels], value=correct_levels[1]['label'], label="Error correction", type="index",
                        interactive=True)

                with gr.Row():
                    promptsTuning = gr.Checkbox(label="Prompts tuning", value=True, interactive=True)
                    anchorStyle = gr.Dropdown(
                        [anchor['label'] for anchor in anchor_styles], value=anchor_styles[0]['label'], label="Anchor style", type="index", interactive=True)

        with gr.Column():
            with gr.Row():
                btn = gr.Button(
                    "Call API",
                    variant="primary"
                )
            with gr.Row():
                out = gr.Image(shape=(1, 1))

    btn.click(greet, [prompt, ControlWeight, errorRate, marginScale, hiresRate, anchorStyle, seed, SizeSelection], out)

demo.launch()