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

from huggingface_hub import whoami

# 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 single_image_generation(prompt: str, image: Optional[str] = None) -> str:
    """Handles text-to-image or single image-to-image."""
    get_fal_key()
    if image:
        image_url = fal_client.upload_file(image)
        result = fal_client.run(
            "fal-ai/nano-banana/edit",
            arguments={"prompt": prompt, "image_url": image_url},
        )
    else:
        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]:
    """Handles multi-image editing."""
    get_fal_key()
    if not images:
        raise gr.Error("Please upload at least one image in the 'Multiple Images' tab.")

    output_images = []
    for image_path in images:
        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()
    pro_message = gr.Markdown(visible=False)
    main_interface = gr.Column(visible=False)

    with main_interface:
        gr.Markdown("## Welcome, PRO User!")
        with gr.Row():
            with gr.Column(scale=1):
                prompt_input = gr.Textbox(
                    label="Prompt",
                    placeholder="A delicious looking pizza"
                )

                with gr.Tabs():
                    with gr.TabItem("Single Image", id="single"):
                        image_input = gr.Image(
                            type="filepath",
                            label="Input Image (Optional for text-to-image)"
                        )
                    with gr.TabItem("Multiple Images", id="multiple"):
                        gallery_input = gr.Gallery(
                            label="Input Images"
                        )

                generate_button = gr.Button("Generate", variant="primary")

            with gr.Column(scale=1):
                # A single gallery can handle one or many images
                output_gallery = gr.Gallery(label="Output")

    def unified_generator(
        prompt: str,
        current_tab: str,
        single_image: Optional[str],
        multi_images: Optional[List[str]],
    ) -> List[str]:
        """
        A single handler that routes to the correct function based on the active tab.
        """
        if current_tab == "multiple":
            return multi_image_edit(prompt, multi_images)
        else: # Handles both text-to-image and single image-to-image
            result_url = single_image_generation(prompt, single_image)
            return [result_url]

    # The `select` event on the tabs gives us the active tab's ID
    selected_tab = gr.Textbox(value="single", visible=False)
    for tab in demo.load_queue[0][0].children[1].children: # Hacky way to get tabs
        if isinstance(tab, gr.Tab):
            tab.select(lambda: tab.id, None, selected_tab)

    generate_button.click(
        unified_generator,
        inputs=[prompt_input, selected_tab, image_input, gallery_input],
        outputs=[output_gallery],
    )

    def control_access(
        profile: Optional[gr.OAuthProfile] = None,
        oauth_token: Optional[gr.OAuthToken] = None
    ):
        """Controls UI visibility based on user's PRO status."""
        if not profile or not oauth_token:
            return gr.update(visible=False), gr.update(visible=False)
        try:
            user_info = whoami(token=oauth_token.token)
            if user_info.get("isPro", False):
                return gr.update(visible=True), gr.update(visible=False)
            else:
                message = (
                    "## ✨ Exclusive Access for PRO Users\n\n"
                    "Thank you for your interest! This feature is available exclusively for our Hugging Face **PRO** members.\n\n"
                    "To unlock this and many other benefits, please consider upgrading your account.\n\n"
                    "### [**Become a PRO Member Today!**](https://huggingface.co/pro)"
                )
                return gr.update(visible=False), gr.update(visible=True, value=message)
        except Exception as e:
            gr.Warning(f"Could not verify user status: {e}")
            return gr.update(visible=False), gr.update(visible=False)

    demo.load(control_access, inputs=None, outputs=[main_interface, pro_message])

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