import gradio as gr
from PIL import Image
import random

from utils import load_learned_embeds, style_guide, IMAGE_SIZE

def generate_output_image(query, selection):
    """Simulates an output image generation based on query and selection."""
    if selection == "None" or not query:
        return None  # No image if conditions aren't met

    # Generating a placeholder image with text (Replace with real logic)
    #img = Image.new("RGB", (512, 512), color=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
    query += ' in the style like BULB'

    img = load_learned_embeds(query, selection)

    return img


def check_visibility(query, selection):
    """Controls visibility of the search button."""
    return gr.update(visible=(selection != "None" and bool(query)))


def ui():
    with gr.Blocks() as demo:
        gr.Markdown("## Image Generation using Stable Diffusion")

        # First Block: Displaying Images with Labels

        with gr.Row():
            gr.Markdown("#### Query: A picture of a mountain landscape during sunset")

        with gr.Row():
            for i in range(6):
                image = f'./examples/{i}.png'
                if i == 0:
                    label = 'Sharpness Loss (Additional Guidance)'
                else:
                    label = style_guide[i-1][0]

                with gr.Column():
                    gr.Image(value=image, label=f"{label}",
                             height=IMAGE_SIZE,
                             width=IMAGE_SIZE,
                             show_download_button=False,
                             show_share_button=False)
                    #gr.Textbox(value=f"Label {label}", interactive=False)

        # Second Block: Query Input and Dropdown
        with gr.Row():
            query_input = gr.Textbox(label="Enter Query")
            dropdown = gr.Dropdown(['None', "Sharpness Loss (Additional Guidance)", 'Oil Style', 'Matrix Style', 'Stripe Style','Dreamy Painting Style','Polygon HD Style'],
                                   label="Select Option")

        # Dynamic Search Button (Initially Hidden)
        search_button = gr.Button("Search", visible=False)
        output_image = gr.Image(label="Output Image", interactive=False)

        # Enable Button Visibility when Conditions are Met
        query_input.change(check_visibility, inputs=[query_input, dropdown], outputs=[search_button])
        dropdown.change(check_visibility, inputs=[query_input, dropdown], outputs=[search_button])

        # On Search, Generate Image
        search_button.click(generate_output_image, inputs=[query_input, dropdown], outputs=[output_image])

    return demo


# Launch the app
demo = ui()
demo.launch()