import io
import random

import gradio as gr
import requests
from PIL import Image

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

def _fetch_url(*args, **kwargs):
    try:
        response = session.get(*args, **kwargs)
        response.raise_for_status()
        return response.content
    except requests.RequestException as e:
        print(f"Request error: {e}")
    return None

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/{prompt}?nologo=true&private=true"
    params = dict(prompt=prompt, seed=seed, width=width, height=height)
    image = _fetch_url(url, params=params)
    if image:
        return Image.open(io.BytesIO(image))
    return None

with gr.Blocks(
    title="Text-to-image UI",
    css="footer {display: none !important}"
) as app:
    gr.Markdown("# Text-to-Image UI")
    with gr.Tab("🎨 Generate"):
        result = gr.Image(
            label="Image",
            show_label=False,
            format="jpeg",
            interactive=False
        )
        with gr.Group():
            with gr.Row():
                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("⚙️ 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
            )

    btn.click(
        generate,
        inputs=[prompt, seed, randomize_seed, width, height],
        outputs=[result],
        api_name="run"
    )

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