import gradio as gr
from huggingface_hub import InferenceClient
import random
import time

# Set up the client for the Hugging Face model
client = InferenceClient(model="GenAIJake/g3na1j8k3")

def generate_image(prompt, negative_prompt="", num_inference_steps=30, guidance_scale=7.5):
    try:
        # Add a small delay to show the loading indicator
        time.sleep(0.5)
        
        # Call the model API
        image = client.text_to_image(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
        )
        
        return image, None
    except Exception as e:
        return None, f"Error generating image: {str(e)}"

# Example prompts to help users get started
example_prompts = [
    ["A serene lake surrounded by mountains at sunset"],
    ["A futuristic cityscape with flying cars"],
    ["A photo-realistic cat wearing sunglasses and a hat"],
    ["A fantasy castle with dragons flying around it"],
    ["An astronaut riding a horse on the moon, digital art"]
]

# Custom CSS for styling
custom_css = """
.gradio-container {
    font-family: 'Arial', sans-serif;
}
.main-header {
    text-align: center;
    margin-bottom: 2rem;
}
.example-prompt {
    cursor: pointer;
    padding: 0.5rem;
    border-radius: 0.5rem;
    background-color: #f7f7f7;
    margin-bottom: 0.5rem;
    transition: background-color 0.3s;
}
.example-prompt:hover {
    background-color: #e0e0e0;
}
"""

# Create the Gradio interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.HTML("""
    <div class="main-header">
        <h1>AI Image Generator</h1>
        <p>Generate amazing images with the power of AI!</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            # Input controls
            with gr.Group():
                prompt = gr.Textbox(
                    label="Enter your prompt",
                    placeholder="Describe the image you want to generate...",
                    lines=3
                )
                
                negative_prompt = gr.Textbox(
                    label="Negative prompt (optional)",
                    placeholder="Elements you want to exclude from the image...",
                    lines=2
                )
                
                with gr.Row():
                    steps = gr.Slider(
                        minimum=10, 
                        maximum=50, 
                        value=30, 
                        step=1, 
                        label="Number of inference steps"
                    )
                    
                    guidance = gr.Slider(
                        minimum=1.0, 
                        maximum=15.0, 
                        value=7.5, 
                        step=0.1, 
                        label="Guidance scale"
                    )
                
                generate_btn = gr.Button("Generate Image", variant="primary")
                
            # Example prompts section
            with gr.Accordion("Example prompts", open=True):
                gr.HTML("<p>Click on any example to use it as your prompt:</p>")
                examples = gr.Examples(
                    example_prompts,
                    inputs=[prompt],
                    examples_per_page=5
                )
        
        with gr.Column(scale=3):
            # Output area with error display
            output_image = gr.Image(label="Generated Image", type="pil")
            error_output = gr.Textbox(label="Error (if any)", visible=False)
    
    # Instructions section
    with gr.Accordion("How to use", open=False):
        gr.HTML("""
        <h3>Instructions:</h3>
        <ul>
            <li><strong>Enter a prompt</strong>: Describe the image you want to generate. Be specific and detailed for better results.</li>
            <li><strong>Negative prompt</strong>: Describe elements you want to avoid in the image.</li>
            <li><strong>Inference steps</strong>: More steps often result in higher quality but take longer to generate.</li>
            <li><strong>Guidance scale</strong>: Higher values make the image adhere more strictly to your prompt.</li>
            <li><strong>Examples</strong>: Click on any example prompt to use it as a starting point.</li>
        </ul>
        """)
    
    # Set up event handlers
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, steps, guidance],
        outputs=[output_image, error_output],
        show_progress=True
    )
    
    # Handle errors by displaying the error message
    error_output.change(
        fn=lambda error: gr.update(visible=error is not None and error != ""),
        inputs=[error_output],
        outputs=[error_output]
    )

# For Hugging Face Spaces, we need to use a specific launch pattern
demo.launch(share=False, server_name="0.0.0.0", server_port=7860)