import random
import gradio as gr
from client import Inference

MAX_SEED = 2 ** 31 - 1
client = Inference()

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)
    return client(
        prompt=prompt,
        seed=seed,
        width=width,
        height=height
    )

with gr.Blocks(
    title="🍒 Cherry Dreamer",
    css="footer {display: none !important}",
    theme=gr.themes.Base(
        primary_hue="red",
        secondary_hue="red",
        neutral_hue="neutral"
    )
) as app:
    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,
                scale=2,
                container=False
            )
            btn = gr.Button("Generate", variant="primary")
            
        with gr.Accordion("⚙️ Settings", open=False):
            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
        )
        clear_btn = gr.Button("Clear")

    with gr.Tab("ℹ️ About"):
        gr.Markdown(f"""
        # 🍒 Cherry Dreamer 
        * Created by [🍒 cherry-ghosts community](https://hf.co/cherry-ghosts)
        * Powered by [🌻 Pollinations.ai](https://pollinations.ai)
        * 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 clear_gallery():
        return []

    clear_btn.click(
        clear_gallery,
        inputs=None,
        outputs=[gallery],
        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(default_concurrency_limit=8).launch(debug=True, ssr_mode=False)