import gradio as gr
from PIL import Image
import asyncio

from text_to_image import TextToImage
from image_to_text import ImageToText
from image_to_image import ImageToImage

# ============================================
# Initialize Model Classes
# ============================================

text_to_image = TextToImage()
image_to_text = ImageToText()
image_to_image = ImageToImage()

# ============================================
# Gradio Interface Functions with Async and Error Handling
# ============================================

async def async_text_to_image(prompt):
    """
    Asynchronous interface function for Text-to-Image generation with error handling.
    """
    try:
        image = await text_to_image.generate_image(prompt)
        return image
    except Exception as e:
        raise gr.Error(f"Text-to-Image Generation Failed: {str(e)}")

async def async_image_to_text(image):
    """
    Asynchronous interface function for Image-to-Text captioning with error handling.
    """
    try:
        caption = await image_to_text.generate_caption(image)
        return caption
    except Exception as e:
        raise gr.Error(f"Image-to-Text Captioning Failed: {str(e)}")

async def async_image_to_image(image, prompt):
    """
    Asynchronous interface function for Image-to-Image transformation with error handling.
    """
    try:
        transformed_image = await image_to_image.transform_image(image, prompt)
        return transformed_image
    except Exception as e:
        raise gr.Error(f"Image-to-Image Transformation Failed: {str(e)}")

# ============================================
# Gradio UI Design
# ============================================

with gr.Blocks(css=".gradio-container {background-color: #f0f8ff}") as demo:
    # Title Section
    gr.Markdown("# 🎨 AI Creativity Hub 🚀")
    gr.Markdown("### Unleash the power of AI to transform your ideas into reality!")

    # Task Selection Radio
    with gr.Tab("✨ Choose Your Magic ✨"):
        task = gr.Radio(
            ["🖼️ Text-to-Image", "📝 Image-to-Text", "🖌️ Image-to-Image"],
            label="Select a Task",
            interactive=True,
            value="🖼️ Text-to-Image"
        )

    # Text-to-Image Section
    with gr.Row(visible=False) as text_to_image_tab:
        with gr.Column():
            gr.Markdown("## 🖼️ Text-to-Image Generator")
            prompt_input = gr.Textbox(
                label="📝 Enter your prompt:",
                placeholder="e.g., A serene sunset over the mountains",
                lines=2
            )
            generate_btn = gr.Button("🎨 Generate Image")
            with gr.Row():
                output_image = gr.Image(label="🖼️ Generated Image")
                download_btn = gr.Button("📥 Download Image")
    
    # Image-to-Text Section
    with gr.Row(visible=False) as image_to_text_tab:
        with gr.Column():
            gr.Markdown("## 📝 Image-to-Text Captioning")
            image_input = gr.Image(
                label="📸 Upload an image:",
                type="pil"
            )
            generate_caption_btn = gr.Button("🖋️ Generate Caption")
            caption_output = gr.Textbox(
                label="📝 Generated Caption:",
                lines=2
            )
    
    # Image-to-Image Section
    with gr.Row(visible=False) as image_to_image_tab:
        with gr.Column():
            gr.Markdown("## 🖌️ Image-to-Image Transformer")
            init_image_input = gr.Image(
                label="📸 Upload an image:",
                type="pil"
            )
            transformation_prompt = gr.Textbox(
                label="📝 Enter transformation prompt:",
                placeholder="e.g., Make it look like a Van Gogh painting",
                lines=2
            )
            transform_btn = gr.Button("🔄 Transform Image")
            with gr.Row():
                transformed_image = gr.Image(label="🖌️ Transformed Image")
                download_transformed_btn = gr.Button("📥 Download Image")
    
    # Define Visibility Based on Task Selection
    def toggle_visibility(selected_task):
        return {
            text_to_image_tab: selected_task == "🖼️ Text-to-Image",
            image_to_text_tab: selected_task == "📝 Image-to-Text",
            image_to_image_tab: selected_task == "🖌️ Image-to-Image",
        }
    
    task.change(
        fn=toggle_visibility, 
        inputs=task, 
        outputs=[text_to_image_tab, image_to_text_tab, image_to_image_tab]
    )

    # Define Button Actions
    generate_btn.click(
        fn=async_text_to_image, 
        inputs=prompt_input, 
        outputs=output_image
    )
    
    download_btn.click(
        fn=lambda img: img.save("generated_image.png") or "Image downloaded!",
        inputs=output_image, 
        outputs=None
    )
    
    generate_caption_btn.click(
        fn=async_image_to_text, 
        inputs=image_input, 
        outputs=caption_output
    )
    
    transform_btn.click(
        fn=async_image_to_image, 
        inputs=[init_image_input, transformation_prompt], 
        outputs=transformed_image
    )
    
    download_transformed_btn.click(
        fn=lambda img: img.save("transformed_image.png") or "Image downloaded!",
        inputs=transformed_image, 
        outputs=None
    )
    
    # Footer Section with Quirky Elements
    gr.Markdown("----")
    gr.Markdown("### 🌟 Explore the endless possibilities with AI! 🌟")
    gr.Markdown("#### 🚀 Built with ❤️ by [Your Name]")

# ============================================
# Launch the Gradio App
# ============================================

demo.launch()