import gradio as gr
from PIL import Image
from models import remix_images as backend_remix_images
from config import OPENAI_API_KEY, GOOGLE_API_KEY, HF_API_TOKEN

# Check if all necessary API keys are set and warn if not
api_keys_missing = []
if not OPENAI_API_KEY:
    api_keys_missing.append("OpenAI (for GPT Image-1)")
if not GOOGLE_API_KEY:
    api_keys_missing.append("Google (for Gemini-2's prompt generation)")
if not HF_API_TOKEN:
    api_keys_missing.append("Hugging Face (for Gemini-2's image generation)")

# Helper function to display missing API key warnings
def display_api_key_warnings():
    if api_keys_missing:
        missing_keys_str = ", ".join(api_keys_missing)
        warning_message = (
            f"🚨 WARNING: The following API keys are not configured as environment variables: {missing_keys_str}. "
            "Some functionalities may be limited or unavailable. Please set them up in your environment for full access."
        )
        print(warning_message)
        gr.Warning(warning_message)
    else:
        print("All required API keys are configured.")


# Gradio Interface
with gr.Blocks(css="#output_image {min-height: 512px;}") as demo:
    gr.Markdown(
        "<center><h1>🌅 Image Remixer 🖼️</h1></center>"
        "Built with <a href='https://huggingface.co/spaces/akhaliq/anycoder' target='_blank'>anycoder</a>"
        "<br>"
        "Drag up to three images into the blank spots, add a text prompt, and select an AI model to remix them into a new image!"
    )

    # Display API key warnings on load
    demo.load(display_api_key_warnings, outputs=None, show_progress="hidden")

    with gr.Row():
        with gr.Column(scale=1):
            image_input1 = gr.Image(type="pil", label="Image Input 1 (Optional)", height=256)
        with gr.Column(scale=1):
            image_input2 = gr.Image(type="pil", label="Image Input 2 (Optional)", height=256)
        with gr.Column(scale=1):
            image_input3 = gr.Image(type="pil", label="Image Input 3 (Optional)", height=256)

    prompt_input = gr.Textbox(
        label="Text Prompt",
        placeholder="A futuristic city blending with ancient ruins, inspired by the input images, vibrant colors.",
        lines=2,
    )

    model_selector = gr.Radio(
        ["GPT Image-1", "Gemini-2"],
        label="Select AI Model for Remixing",
        value="GPT Image-1", # Default value
        info="GPT Image-1 uses DALL-E 3 (prompted by Gemini Pro Vision). Gemini-2 uses Google Gemini Pro Vision to generate a detailed prompt, which is then fed to Stable Diffusion XL.",
    )

    submit_button = gr.Button("Remix Images", variant="primary")

    output_image = gr.Image(label="Remixed Output", elem_id="output_image")

    submit_button.click(
        fn=backend_remix_images,
        inputs=[model_selector, image_input1, image_input2, image_input3, prompt_input],
        outputs=output_image,
        api_name="remix",
    )

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