import os
import gradio as gr
import numpy as np
from config.options import aspect_ratio
from config.theme_config import theme

css = """
#sidebar {
    background-color: #27272A;
    padding: 10px;
}
.ml {
    padding-left: 10px;
    padding-top: 4px;
}
.gr-group {
    background-color: red!important;
}
"""

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))


def greet(name):
    return "Hello " + name + "!"


with gr.Blocks(theme=theme, css=css) as demo:
    with gr.Row():
        with gr.Column(scale=1, elem_id="sidebar"):
            title = gr.Markdown("Prompt")
            prompt = gr.TextArea(
                label="Prompt",
                placeholder="Enter your prompt",
                lines=3,
                show_label=False,
                container=False,
            )

            with gr.Column():
                gr.Markdown("Aspect Ratio")
                gallery = gr.Gallery(
                    [(item["image"], item["name"]) for item in aspect_ratio],
                    label="Gallery",
                    show_label=False,
                    allow_preview=False,
                    elem_id="gallery",
                    show_share_button=False,
                    # height=134,
                    columns=3,
                    container=False,
                    object_fit="scale-down",
                )

            gr.Markdown("Advanced Settings")
            with gr.Accordion(
                "Show/Hide",
                open=True,
            ):
                with gr.Group():
                    gr.Markdown(
                        "Seed:",
                        container=False,
                        elem_classes="ml",
                    )
                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=0,
                        container=False,
                    )
                # randomize_seed = gr.Checkbox(label="Randomize seed", value=True, info="If checked, a random seed will be used for each run.")

                with gr.Group():
                    gr.Markdown(
                        "Width:",
                        container=False,
                        elem_classes="ml",
                    )
                    width = gr.Slider(
                        label="Width",
                        minimum=256,
                        maximum=MAX_IMAGE_SIZE,
                        step=32,
                        value=1024,  # Replace with defaults that work for your model
                        container=False,
                    )

                with gr.Group():
                    gr.Markdown("Height:", container=False, elem_classes="ml")
                    height = gr.Slider(
                        label="Height",
                        minimum=256,
                        maximum=MAX_IMAGE_SIZE,
                        step=32,
                        value=1024,  # Replace with defaults that work for your model
                        container=False,
                    )

                with gr.Group():
                    gr.Markdown(
                        "Guidance Strength:", container=False, elem_classes="ml"
                    )
                    guidance_scale = gr.Slider(
                        label="Guidance scale",
                        minimum=0.0,
                        maximum=10.0,
                        step=0.1,
                        value=6.0,  # Replace with defaults that work for your model
                        container=False,
                    )

                with gr.Group():
                    gr.Markdown(
                        "Number of Inference Steps:", container=False, elem_classes="ml"
                    )
                    num_inference_steps = gr.Slider(
                        label="Number of inference steps",
                        minimum=1,
                        maximum=50,
                        step=1,
                        value=25,  # Replace with defaults that work for your model
                        container=False,
                    )
            run_btn = gr.Button(
                "Run",
                variant="primary",
            )

        with gr.Column(scale=2):
            output = gr.Image(label="Output Box", show_label=False)

    run_btn.click(fn=greet, inputs=prompt, outputs=output, api_name="greet")

if __name__ == "__main__":
    demo.launch()