import random
import time
from io import BytesIO
from urllib.parse import quote

from datasets import load_dataset
import gradio as gr
import requests
from PIL import Image

MAX_SEED = 2 ** 31 - 1
session = requests.Session()

def _fetch(*args, retries=5, **kwargs):
    for i in range(retries + 1):
        try:
            response = session.get(*args, **kwargs)
            response.raise_for_status()
            return response.content
        except requests.RequestException:
            if response.status_code in {429, 500, 502, 503, 504}:
                delay = min(2 ** i, 60)
                time.sleep(delay)
            else:
                return None
    return None

def load_examples():
    ds = load_dataset("k-mktr/improved-flux-prompts", split="train")
    return [[item] for item in ds.shuffle()["prompt"][:20]]

def generate(
    prompt,
    seed=None,
    randomize_seed=True,
    width=1024,
    height=1024
):
    if seed is None or randomize_seed:
        seed = random.randint(0, MAX_SEED)
    url = f"https://pollinations.ai/p/{quote(prompt)}?nologo=true&private=true"
    params = dict(
        prompt=prompt,
        seed=seed,
        width=width,
        height=height
    )
    image = _fetch(url, params=params)
    if image:
        return Image.open(BytesIO(image))
    return None

with gr.Blocks(
    title="Text-to-Image",
    css="footer {display: none !important}",
    theme=gr.themes.Base(
        primary_hue="red",
        secondary_hue="red",
        neutral_hue="neutral"
    )
) as app:
    gr.Markdown("# Text-to-Image")
    with gr.Tab("🎨 Generate"):
        result = gr.Image(
            label="Image",
            show_label=False,
            format="jpeg",
            interactive=False
        )
        with gr.Row(equal_height=True):
            prompt = gr.Textbox(
                label="Prompt",
                show_label=False,
                placeholder="Enter your prompt..",
                max_lines=1,
                container=False
            )
            btn = gr.Button("Generate!", variant="primary")

    with gr.Tab("💡 Examples"):
        examples = gr.Examples(
            load_examples(),
            inputs=[prompt],
            examples_per_page=2
        )
        randomize_btn = gr.Button("🎲 Randomize")
            
    with gr.Tab("⚙️ Settings"):
        with gr.Row():
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

        with gr.Row():
            width = gr.Slider(
                label="Width",
                minimum=256,
                maximum=1344,
                step=64,
                value=1024
            )
            height = gr.Slider(
                label="Height",
                minimum=256,
                maximum=1344,
                step=64,
                value=1024
            )

    with gr.Tab("🖼️ Gallery"):
        gallery = gr.Gallery(
            label="Gallery",
            show_label=False,
            format="jpeg",
            interactive=False
        )

    with gr.Accordion("ℹ️ About", open=False):
        gr.Markdown(f"""
        * Created by [🍒 cherry-ghosts community](https://hf.co/cherry-ghosts)
        * Powered by [🐝 Pollinations](https://pollinations.ai) | [GitHub](https://github.com/pollinations/pollinations)
        * Running on [Gradio](https://www.gradio.app) v{gr.__version__}
        """)

    def add_to_gallery(img, gallery):
        if gallery is None:
            gallery = []
        if img is not None:
            gallery = gallery + [img]
        return gallery

    def update_examples():
        return gr.update(samples=load_examples())

    app.load(
        update_examples,
        inputs=None,
        outputs=[examples.dataset],
        queue=False,
        show_api=False
    )

    randomize_btn.click(
        update_examples,
        inputs=None,
        outputs=[examples.dataset],
        queue=False,
        show_api=False
    )

    btn.click(
        generate,
        inputs=[prompt, seed, randomize_seed, width, height],
        outputs=[result],
        api_name="run"
    ).then(
        add_to_gallery,
        inputs=[result, gallery],
        outputs=[gallery],
        queue=False,
        show_api=False
    )

if __name__ == "__main__":
    app.queue().launch(debug=True, ssr_mode=False)