import gradio as gr
from PIL import Image
from models import gemini_remix_image, gpt_remix_image

def remix_images_wrapper(image1_path: str | None, image2_path: str | None, image3_path: str | None, prompt: str) -> tuple[Image.Image, Image.Image]:
    """
    Wrapper function to process input paths, convert to PIL Images,
    and call both mock remixing models.
    """
    pil_images = []
    for path in [image1_path, image2_path, image3_path]:
        if path is not None:
            pil_images.append(Image.open(path))
        else:
            pil_images.append(None) # Maintain None for absent images

    # Call the mock model functions
    gemini_output = gemini_remix_image(pil_images[0], pil_images[1], pil_images[2], prompt)
    gpt_output = gpt_remix_image(pil_images[0], pil_images[1], pil_images[2], prompt)

    return gemini_output, gpt_output

with gr.Blocks(title="AI Image Remixer") as demo:
    gr.Markdown(
        "<div style='text-align: center; margin-bottom: 20px;'>"
        "<h1>🎨 AI Image Remixer 🖼️</h1>"
        "<p>Drag up to three images into the input boxes, add a text prompt, and see them remixed by different AI models!</p>"
        "</div>"
    )
    gr.Markdown(
        "<div style='text-align: center; margin-bottom: 20px;'>"
        "<p>Built with <a href='https://huggingface.co/spaces/akhaliq/anycoder' target='_blank'>anycoder</a></p>"
        "</div>"
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("<h3>Input Images (Drag & Drop)</h3>")
            with gr.Row():
                input_image1 = gr.Image(
                    type="filepath", label="Image 1", 
                    height=200, width=200, 
                    sources=["upload", "clipboard"],
                    value="assets/example1.png" # Example image 1
                )
                input_image2 = gr.Image(
                    type="filepath", label="Image 2", 
                    height=200, width=200, 
                    sources=["upload", "clipboard"],
                    value="assets/example2.png" # Example image 2
                )
                input_image3 = gr.Image(
                    type="filepath", label="Image 3", 
                    height=200, width=200, 
                    sources=["upload", "clipboard"],
                    value="assets/example3.png" # Example image 3
                )
            
            prompt_textbox = gr.Textbox(
                label="Text Prompt", 
                placeholder="Describe how to remix the images, e.g., 'futuristic city, vibrant colors'", 
                lines=2,
                value="Create a magical forest scene with bioluminescent plants."
            )
            remix_button = gr.Button("Remix Images!")
        
        with gr.Column(scale=2):
            gr.Markdown("<h3>Remixed Outputs</h3>")
            with gr.Row():
                gemini_output_image = gr.Image(label="Gemini-2 Remix", height=400, show_share_button=True)
                gpt_output_image = gr.Image(label="GPT Image-1 Remix", height=400, show_share_button=True)
    
    remix_button.click(
        fn=remix_images_wrapper,
        inputs=[input_image1, input_image2, input_image3, prompt_textbox],
        outputs=[gemini_output_image, gpt_output_image],
        api_name="remix_images",
        show_progress="minimal"
    )

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