import os
print("Starting AI Headshot Generator...")

import gradio as gr
import numpy as np
import spaces  # ZeroGPU decorator comes from this package, not huggingface_hub

# ======================
# Simple license system
# ======================
ADMIN_KEYS = {"HEADSHOT-TEST123456", "HEADSHOT-OWNERACCESS"}
MAX_TRIALS = 3


def verify_license(key):
    if not key:
        return False
    return key.strip().upper() in ADMIN_KEYS


# ======================
# ZeroGPU-decorated function
# ======================
@spaces.GPU(duration=180)  # 180 seconds max per generation
def generate_headshot(image, license_key):
    """
    This function runs on a real GPU when executed in a ZeroGPU Space.
    Replace the placeholder logic later with your real InstantID / IP-Adapter code.
    """
    # ------------------ License check ------------------
    if not verify_license(license_key):
        blank = np.zeros((512, 512, 3), dtype=np.uint8)
        return blank, "Invalid license. Use: HEADSHOT-TEST123456"

    # ------------------ Real inference goes here ------------------
    if image is None:
        h, w = 512, 512
    else:
        h, w = image.shape[:2]

    # Light blue/gray placeholder so you instantly see it works
    result = np.ones((h, w, 3), dtype=np.uint8) * 200
    # Pink rectangle just for fun
    result[50:150, 50:450] = [255, 182, 193]

    return result, "Generated! GPU is working — ready for InstantID pipeline!"


# ======================
# Gradio Interface
# ======================
with gr.Blocks(title="AI Headshot Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # AI Headshot Generator
        *Transform selfies into professional headshots*
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Upload Your Photo")
            image_input = gr.Image(label="", type="numpy", height=300)

            gr.Markdown("### License Key")
            license_input = gr.Textbox(
                label="",
                value="HEADSHOT-TEST123456",
                placeholder="Enter HEADSHOT-TEST123456",
            )

            gr.Markdown(
                f"""
                ### How It Works
                1. Upload a photo  
                2. Use test key: **HEADSHOT-TEST123456**  
                3. Get {MAX_TRIALS} free generations  
                4. Purchase license → HD + multiple styles  
                [Purchase License](https://canadianheadshotpro.carrd.co)
                """
            )

            generate_btn = gr.Button("Generate Headshot", variant="primary", size="lg")
            status_output = gr.Textbox(label="Status", interactive=False)

        with gr.Column(scale=1):
            gr.Markdown("### Your Headshot")
            result_image = gr.Image(label="", height=500)

            gr.Markdown(
                """
                ### Download Options
                - **Free Trial**: Basic preview  
                - **Premium**: HD quality  
                - **Professional**: Multiple styles  

                AI models loading soon!
                """
            )

    # Connect button
    generate_btn.click(
        fn=generate_headshot,
        inputs=[image_input, license_input],
        outputs=[result_image, status_output],
    )

    # Footer
    gr.Markdown(
        """
        ---
        <div style="text-align: center; color: #888; margin-top: 20px;">
        Powered by Hugging Face ZeroGPU • InstantID coming soon • © 2025 Bipin Pandey
        </div>
        """
    )


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