import gradio as gr
from models import generate_image, MODEL_ID
from config import APPLE_TENCENT_THEME
from typing import Optional, Union
from huggingface_hub import whoami

def verify_pro_status(token: Optional[Union[gr.OAuthToken, str]]) -> bool:
    """Verifies if the user is a Hugging Face PRO user or part of an enterprise org."""
    if not token:
        return False
    
    if isinstance(token, gr.OAuthToken):
        token_str = token.token
    elif isinstance(token, str):
        token_str = token
    else:
        return False
    
    try:
        user_info = whoami(token=token_str)
        return (
            user_info.get("isPro", False) or
            any(org.get("isEnterprise", False) for org in user_info.get("orgs", []))
        )
    except Exception as e:
        print(f"Could not verify user's PRO/Enterprise status: {e}")
        return False

def generate_image_with_pro_check(prompt: str, profile: Optional[gr.OAuthProfile] = None, oauth_token: Optional[gr.OAuthToken] = None):
    """
    Wrapper function that checks if user is PRO before generating image.
    """
    if profile is None:
        raise gr.Error("Please Sign in with Hugging Face to use this app")
    
    # Check if user has PRO status
    if not verify_pro_status(oauth_token):
        raise gr.Error(
            "This app is exclusively for Hugging Face PRO subscribers. "
            "Please upgrade at https://huggingface.co/subscribe/pro to access this feature."
        )
    
    # User is PRO, proceed with image generation
    return generate_image(prompt)

def create_ui():
    with gr.Blocks(title=f"Tencent HunyuanImage-3.0 Demo", theme=APPLE_TENCENT_THEME) as demo:
        gr.HTML(
            f"<div style='text-align: center; max-width: 700px; margin: 0 auto;'>"
            f"<h1>Tencent {MODEL_ID.split('/')[-1]}</h1>"
            f"<p>Generate images using Tencent's state-of-the-art model hosted by FAL AI.</p>"
            f"<p style='color: orange;'>⚠️ This app is exclusively for Hugging Face PRO subscribers.</p>"
            f"<p><a href='https://huggingface.co/subscribe/pro' target='_blank' style='color: #ff6b6b; font-weight: bold;'>🚀 Subscribe to PRO</a></p>"
            f"Built with <a href='https://huggingface.co/spaces/akhaliq/anycoder' target='_blank'>anycoder</a>"
            f"</div>"
        )
        
        # Add login button - required for OAuth
        gr.LoginButton()
        
        # PRO status message area
        pro_message = gr.Markdown(visible=False)
        main_interface = gr.Column(visible=False)

        with main_interface:
            with gr.Row():
                with gr.Column(scale=1):
                    prompt_input = gr.Textbox(
                        label="Prompt",
                        placeholder="e.g., A detailed watercolor painting of a small red fox sleeping on a pile of autumn leaves.",
                        lines=4
                    )
                    generate_btn = gr.Button("🎨 Generate Image", variant="primary")
                    
                with gr.Column(scale=1):
                    output_image = gr.Image(
                        label="Generated Image",
                        height=512,
                        width=512,
                        interactive=False,
                        show_download_button=True
                    )
            
            # Set up the event listener - note the function now takes OAuthProfile and OAuthToken
            generate_btn.click(
                fn=generate_image_with_pro_check,
                inputs=[prompt_input],
                outputs=[output_image],
                queue=False,
                api_name=False,
                show_api=False,
            )
            
            # Example usage guidance with queue features disabled
            gr.Examples(
                examples=[
                    "A detailed watercolor painting of a small red fox sleeping on a pile of autumn leaves.",
                    "A futuristic city skyline at sunset with flying cars",
                    "A mystical forest with glowing mushrooms and fireflies",
                    "An astronaut riding a horse on Mars",
                ],
                inputs=prompt_input,
                outputs=output_image,
                fn=generate_image_with_pro_check,
                cache_examples=False,
                api_name=False,
                show_api=False,
            )
            
            gr.Markdown("<h3 style='text-align: center'>Thank you for being a PRO subscriber! 🤗</h3>")

        # Control access based on PRO status
        def control_access(profile: Optional[gr.OAuthProfile] = None, oauth_token: Optional[gr.OAuthToken] = None):
            if not profile:
                return (
                    gr.update(visible=False),  # main_interface
                    gr.update(visible=True, value=(
                        "## 👋 Welcome!\n\n"
                        "Please **Sign in with Hugging Face** using the button above to continue."
                    ))  # pro_message
                )
            
            if verify_pro_status(oauth_token):
                return (
                    gr.update(visible=True),   # main_interface
                    gr.update(visible=False)   # pro_message
                )
            else:
                message = (
                    "## ✨ Exclusive Access for PRO Users\n\n"
                    "Thank you for your interest! This app is available exclusively for our Hugging Face **PRO** members.\n\n"
                    "With PRO, you'll get:\n"
                    "- Access to exclusive Spaces like this one\n"
                    "- Increased GPU quotas\n"
                    "- Early access to new features\n"
                    "- Priority support\n\n"
                    "### [**Become a PRO Today!**](https://huggingface.co/subscribe/pro)\n\n"
                    "*Already a PRO member? Try logging out and back in if you're having issues.*"
                )
                return (
                    gr.update(visible=False),  # main_interface
                    gr.update(visible=True, value=message)  # pro_message
                )
        
        # Load event to check access on page load
        demo.load(control_access, inputs=None, outputs=[main_interface, pro_message])

    return demo

if __name__ == "__main__":
    app = create_ui()
    # Launch without special auth parameters
    # OAuth is enabled via Space metadata (hf_oauth: true in README.md)
    app.launch(
        show_api=False,
        enable_monitoring=False,
        quiet=True,
    )