import gradio as gr
import fal_client
import os
from typing import Optional

# It is recommended to create this as a Secret on your Hugging Face Space
# For example: FAL_KEY = "fal_key_..."
FAL_KEY = os.getenv("FAL_KEY", "")

# Set the key for the fal_client
if FAL_KEY:
    fal_client.api_key = FAL_KEY

def get_fal_key():
    """Checks for the FAL_KEY and raises a Gradio error if it's not set."""
    if not FAL_KEY:
        raise gr.Error("FAL_KEY is not set. Please add it to your Hugging Face Space secrets.")

def generate_image(prompt: str, image: Optional[str] = None) -> str:
    """
    Generates or edits an image.
    - If an image filepath is provided, it performs image-to-image editing.
    - Otherwise, it performs text-to-image generation.
    """
    get_fal_key()

    if image:
        # Image-to-Image: Upload the local file to get a public URL
        image_url = fal_client.upload_file(image)
        result = fal_client.run(
            "fal-ai/nano-banana/edit",
            arguments={
                "prompt": prompt,
                "image_url": image_url,
            },
        )
    else:
        # Text-to-Image
        result = fal_client.run(
            "fal-ai/nano-banana",
            arguments={
                "prompt": prompt,
            },
        )
    return result["images"][0]["url"]

def multi_image_edit(prompt: str, images: list[str]) -> list[str]:
    """Edits multiple images based on a text prompt."""
    get_fal_key()

    if not images:
        raise gr.Error("Please upload at least one image.")

    output_images = []
    for image_path in images:
        # Upload each image and get its URL
        image_url = fal_client.upload_file(image_path)
        result = fal_client.run(
            "fal-ai/nano-banana/edit",
            arguments={
                "prompt": prompt,
                "image_url": image_url,
            },
        )
        output_images.append(result["images"][0]["url"])

    return output_images

# --- Gradio App UI ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Nano Banana Image Generation")
    gr.Markdown("Generate or edit images with FAL. **Sign in with Hugging Face to begin.**")
    
    login_button = gr.LoginButton()
    
    # This entire Column will be hidden until the user logs in
    main_interface = gr.Column(visible=False)

    with main_interface:
        gr.Markdown("## Welcome! You are logged in.")
        with gr.Tabs():
            with gr.TabItem("Generate"):
                with gr.Row():
                    with gr.Column(scale=1):
                        prompt_input = gr.Textbox(label="Prompt", placeholder="A delicious looking pizza")
                        image_input = gr.Image(type="filepath", label="Input Image (Optional)")
                        generate_button = gr.Button("Generate", variant="primary")
                    with gr.Column(scale=1):
                        generate_output = gr.Image(label="Output")
                
                generate_button.click(
                    generate_image,
                    inputs=[prompt_input, image_input],
                    outputs=[generate_output],
                )

            with gr.TabItem("Multi-Image Edit"):
                with gr.Row():
                    with gr.Column(scale=1):
                        multi_prompt_input = gr.Textbox(label="Prompt", placeholder="Make it black and white")
                        multi_image_input = gr.Gallery(label="Input Images", file_types=["image"])
                        multi_edit_button = gr.Button("Edit Images", variant="primary")
                    with gr.Column(scale=1):
                        multi_image_output = gr.Gallery(label="Output Images")
                
                multi_edit_button.click(
                    multi_image_edit,
                    inputs=[multi_prompt_input, multi_image_input],
                    outputs=[multi_image_output],
                )

    def show_app_on_login(profile: Optional[gr.OAuthProfile] = None):
        """
        Controls the visibility of the main UI.
        Gradio automatically injects `gr.OAuthProfile` if the user is logged in.
        """
        if profile:
            return gr.update(visible=True)
        return gr.update(visible=False)

    # When the app loads or the user logs in/out, this function will be called.
    # It receives the user's profile (or None) and updates the UI visibility.
    demo.load(show_app_on_login, inputs=None, outputs=main_interface)

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