import gradio as gr
import os
from datetime import datetime

def generate_image(prompt, style, num_images=1, width=512, height=512):
    """
    Generate images based on the given prompt and parameters.
    Returns list of generated image paths.
    """
    try:
        model = gr.load("models/LAYEK-143/FLUX_V0")
        images = []
        for i in range(num_images):
            # Generate unique filename with timestamp
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"generated_{timestamp}_{i}.png"
            
            # Generate the image
            result = model.predict(
                prompt=prompt,
                style=style,
                width=width,
                height=height
            )
            
            # Save the image
            os.makedirs("generated_images", exist_ok=True)
            save_path = os.path.join("generated_images", filename)
            result.save(save_path)
            images.append(save_path)
            
        return images
    except Exception as e:
        return str(e)

def create_interface():
    """Create and configure the Gradio interface"""
    
    # Define available style options
    STYLES = [
        "Realistic",
        "Artistic",
        "Abstract",
        "Cartoon",
        "Sketch"
    ]
    
    # Create the interface
    with gr.Blocks(title="Image Generator") as interface:
        gr.Markdown("# 🎨 Image Generator")
        
        with gr.Row():
            with gr.Column():
                # Input components
                prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Describe the image you want to generate...",
                    lines=3
                )
                
                style = gr.Dropdown(
                    choices=STYLES,
                    label="Style",
                    value="Realistic"
                )
                
                with gr.Row():
                    num_images = gr.Slider(
                        minimum=1,
                        maximum=4,
                        value=1,
                        step=1,
                        label="Number of Images"
                    )
                    
                with gr.Row():
                    width = gr.Slider(
                        minimum=256,
                        maximum=1024,
                        value=512,
                        step=64,
                        label="Width"
                    )
                    height = gr.Slider(
                        minimum=256,
                        maximum=1024,
                        value=512,
                        step=64,
                        label="Height"
                    )
                
                generate_btn = gr.Button("Generate Images", variant="primary")
            
            with gr.Column():
                # Output gallery
                output_gallery = gr.Gallery(
                    label="Generated Images",
                    show_label=True,
                    elem_id="gallery"
                ).style(grid=2, height="auto")
        
        # Error message display
        error_message = gr.Textbox(
            label="Status",
            visible=False
        )
        
        # Handle generation
        def on_generate(prompt, style, num_images, width, height):
            if not prompt.strip():
                return None, "Please enter a prompt."
            
            try:
                images = generate_image(prompt, style, num_images, width, height)
                if isinstance(images, str):  # Error occurred
                    return None, images
                return images, "Generation successful!"
            except Exception as e:
                return None, f"Error: {str(e)}"
        
        generate_btn.click(
            fn=on_generate,
            inputs=[prompt, style, num_images, width, height],
            outputs=[output_gallery, error_message]
        )
        
        # Add usage instructions
        gr.Markdown("""
        ## How to Use
        1. Enter a detailed description of the image you want to generate
        2. Select a style from the dropdown menu
        3. Adjust the number of images and dimensions as needed
        4. Click 'Generate Images' and wait for the results
        5. Generated images will be saved in the 'generated_images' folder
        
        Note: Generation time may vary depending on the complexity of your prompt.
        """)
    
    return interface

if __name__ == "__main__":
    # Create and launch the interface
    app = create_interface()
    app.launch(
        share=True,
        enable_queue=True,
        server_name="0.0.0.0",
        server_port=7860
    )