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 verify_pro_status(token: Optional[gr.OAuthToken]) -> bool:
    """Verifies if the user is a Hugging Face PRO user using their token."""
    if not token:
        return False
    try:
        user_info = whoami(token=token.token)
        return user_info.get("isPro", False)
    except Exception as e:
        print(f"Could not verify user's PRO status: {e}")
        return False

# --- Backend Generation Functions ---

def run_single_image_logic(prompt: str, image: Optional[str] = None) -> List[str]:
    """Handles text-to-image or single image-to-image and returns a list."""
    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 run_multi_image_logic(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 for PROs")
    gr.Markdown("Hugging Face PRO users can use Google's Nano Banana (Gemini 2.5 Flash Image Preview) on this Space. [Subscribe to PRO](https://huggingface.co/pro)")

    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():
            # LEFT COLUMN: Inputs
            with gr.Column(scale=1):
                prompt_input = gr.Textbox(
                    label="Prompt",
                    placeholder="A delicious looking pizza"
                )
                active_tab_state = gr.State(value="single")
                with gr.Tabs() as tabs:
                    with gr.TabItem("Single Image", id="single") as single_tab:
                        image_input = gr.Image(
                            type="filepath",
                            label="Input Image (Optional for text-to-image)"
                        )
                    with gr.TabItem("Multiple Images", id="multiple") as multi_tab:
                        gallery_input = gr.Gallery(
                            label="Input Images", file_types=["image"]
                        )
                generate_button = gr.Button("Generate", variant="primary")

            # RIGHT COLUMN: Outputs
            with gr.Column(scale=1):
                output_gallery = gr.Gallery(label="Output")
                selected_output_image_state = gr.State()
                use_image_button = gr.Button("♻️ Use Generated Image for Next Edit")

    # --- Event Handlers ---

    def unified_generator(
        prompt: str,
        single_image: Optional[str],
        multi_images: Optional[List[str]],
        active_tab: str,
        oauth_token: Optional[gr.OAuthToken] = None,
    ) -> List[str]:
        if not verify_pro_status(oauth_token):
            raise gr.Error("Access Denied. This service is for PRO users only.")
        if active_tab == "multiple" and multi_images:
            return run_multi_image_logic(prompt, multi_images)
        else:
            return run_single_image_logic(prompt, single_image)

    single_tab.select(lambda: "single", None, active_tab_state)
    multi_tab.select(lambda: "multiple", None, active_tab_state)
    
    generate_button.click(
        unified_generator,
        inputs=[prompt_input, image_input, gallery_input, active_tab_state, login_button],
        outputs=[output_gallery],
    )

    # New handlers for the continuous editing loop
    def store_selected_image(evt: gr.SelectData):
        """When an image is selected in the output gallery, store its path in state."""
        return evt.value['image']

    def reuse_output_image(selected_image_path):
        """
        Takes the path from state and sends it to the single image input.
        Also forces the UI to switch to the "Single Image" tab.
        """
        if not selected_image_path:
            gr.Warning("Please select an image from the output gallery first!")
            return None, gr.update()
        # Output 1: The image path for the gr.Image component
        # Output 2: The ID of the tab to select for the gr.Tabs component
        return selected_image_path, "single"

    output_gallery.select(store_selected_image, None, selected_output_image_state)
    use_image_button.click(
        reuse_output_image,
        inputs=[selected_output_image_state],
        outputs=[image_input, tabs]
    )

    # --- Access Control Logic ---
    def control_access(
        profile: Optional[gr.OAuthProfile] = None,
        oauth_token: Optional[gr.OAuthToken] = None
    ):
        if not profile:
            return gr.update(visible=False), gr.update(visible=False)
        if verify_pro_status(oauth_token):
            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)

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

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