import gradio as gr
from utils import generate_images_simulated

# --- Constants ---
ANYCODER_LINK = "https://huggingface.co/spaces/akhaliq/anycoder"
MODEL_CHOICES = {
    "Stable Diffusion v1.4 (CompVis)": "CompVis/stable-diffusion-v1-4",
    "Stable Diffusion 2.1 Base (Manojb)": "Manojb/stable-diffusion-2-1-base"
}
DEFAULT_MODEL_KEY = list(MODEL_CHOICES.keys())[0]

# --- Gradio Logic Wrapper ---
def process_generation(prompt: str, negative_prompt: str, model_key: str, steps: float, guidance_scale: float, num_images: float):
    """
    Wrapper function to call the simulated generation logic, translating model key.
    
    Args:
        prompt: The text prompt for generation.
        negative_prompt: Text prompt to guide the model away from.
        model_key: The user-friendly name of the selected model.
        steps: Inference steps.
        guidance_scale: CFG scale.
        num_images: Number of images to generate (1-4).
        
    Returns:
        A list of PIL Images for the Gradio Gallery.
    """
    model_name_id = MODEL_CHOICES[model_key]
    
    # Note: steps and guidance_scale are floats from the slider, cast to appropriate types
    results = generate_images_simulated(
        prompt=prompt, 
        negative_prompt=negative_prompt, 
        model_name=model_name_id, 
        steps=int(steps), 
        guidance_scale=guidance_scale, 
        num_images=int(num_images)
    )
    return results

# --- Gradio UI Definition ---
with gr.Blocks(title="Stable Diffusion Model Comparison Demo") as demo:
    
    # Custom Header
    gr.HTML(
        f"""
        <div style="text-align: center; max-width: 800px; margin: 0 auto;">
            <h1>Stable Diffusion Model Comparison Demo</h1>
            <p>Compare outputs from Stable Diffusion v1.4 and 2.1 Base. The application handles empty prompts by using a default prompt internally.</p>
            <p>Built with <a href="{ANYCODER_LINK}" target="_blank" style="color: blue; text-decoration: underline;">anycoder</a></p>
        </div>
        """
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            # --- Inputs ---
            prompt = gr.Textbox(
                label="Prompt", 
                placeholder="e.g., A photorealistic image of a vintage robot reading a book in a cozy library, cinematic lighting",
                value="A majestic portrait of a futuristic cyborg cat, digital art, highly detailed",
                lines=3
            )
            
            negative_prompt = gr.Textbox(
                label="Negative Prompt (Optional)", 
                placeholder="e.g., blurry, low quality, distorted, bad anatomy, text",
                lines=2
            )
            
            model_selector = gr.Radio(
                label="Select Model",
                choices=list(MODEL_CHOICES.keys()),
                value=DEFAULT_MODEL_KEY,
                interactive=True
            )
            
            with gr.Accordion("Generation Parameters", open=True):
                steps = gr.Slider(
                    minimum=10, maximum=100, step=5, value=50, label="Inference Steps"
                )
                
                guidance_scale = gr.Slider(
                    minimum=1.0, maximum=15.0, step=0.5, value=7.5, label="Guidance Scale (CFG)"
                )
                
                num_images = gr.Slider(
                    minimum=1, maximum=4, step=1, value=4, label="Number of Images to Generate",
                    info="Uploads up to 4 images."
                )
            
            generate_btn = gr.Button("Generate Images", variant="primary")

        with gr.Column(scale=3):
            # --- Outputs ---
            output_gallery = gr.Gallery(
                label="Generated Images (512x512)", 
                columns=2, 
                rows=2, 
                height="auto", 
                preview=True
            )
            
            gr.Examples(
                examples=[
                    ["A hyper-detailed oil painting of a lighthouse during a storm", "ugly, blurry, low resolution", DEFAULT_MODEL_KEY],
                    ["", "text, watermark, poor rendering", "Stable Diffusion 2.1 Base (Manojb)"], # Empty prompt example
                    ["A 1950s sci-fi movie poster featuring a giant insect attacking Paris", "modern, photorealistic, color photo", DEFAULT_MODEL_KEY]
                ],
                inputs=[prompt, negative_prompt, model_selector],
                label="Examples (Note the second example uses an empty prompt)"
            )

    # --- Logic Binding ---
    generate_btn.click(
        fn=process_generation,
        inputs=[prompt, negative_prompt, model_selector, steps, guidance_scale, num_images],
        outputs=[output_gallery]
    )
    
    # Run initial generation on load to populate the gallery immediately
    demo.load(
        fn=process_generation,
        inputs=[prompt, negative_prompt, model_selector, steps, guidance_scale, num_images],
        outputs=[output_gallery]
    )

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