import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image


def generate_image(api_key, prompt, model_choice):
    """
    Generate image using Hugging Face Inference API
    """

    # Validation
    if not api_key.strip():
        raise gr.Error("Please enter your Hugging Face API Key")

    if not prompt.strip():
        raise gr.Error("Please enter an image prompt")

    try:
        # Determine model and provider based on model_choice
        if model_choice == "Stable Diffusion XL (Free)":
            model = "stabilityai/stable-diffusion-xl-base-1.0"
            provider = None
        elif model_choice == "FLUX.1-schnell (Free)":
            model = "black-forest-labs/FLUX.1-schnell"
            provider = None
        elif model_choice == "FLUX.1-dev (Premium)":
            model = "black-forest-labs/FLUX.1-dev"
            provider = "fal-ai"
        else:
            raise gr.Error("Unknown model selected")

        # Create HF client
        if provider:
            client = InferenceClient(
                provider=provider,
                api_key=api_key
            )
        else:
            client = InferenceClient(
                api_key=api_key
            )

        # Generate image
        image = client.text_to_image(
            prompt,
            model=model
        )

        return image

    except Exception as e:
        raise gr.Error(f"Generation Failed: {str(e)}")


# Custom CSS
css = """
#generate_btn {
    height: 60px;
    font-size: 18px;
    font-weight: bold;
}
"""


with gr.Blocks(
    title="AI Image Generator"
) as demo:

    gr.Markdown(
        """
        # 🖼️ AI Image Generator
        Generate images using Hugging Face Text-to-Image models (Free & Premium)
        """
    )

    with gr.Row():

        # LEFT PANEL
        with gr.Column(scale=1):

            api_key = gr.Textbox(
                label="Hugging Face API Key",
                placeholder="hf_xxxxxxxxxxxxxxxxxxxxxxxxx",
                type="password"
            )

            model_choice = gr.Dropdown(
                label="Model Selection",
                choices=[
                    "Stable Diffusion XL (Free)",
                    "FLUX.1-schnell (Free)",
                    "FLUX.1-dev (Premium)"
                ],
                value="Stable Diffusion XL (Free)"
            )

            prompt = gr.Textbox(
                label="Image Prompt",
                placeholder="Astronaut riding a horse on Mars, cinematic lighting...",
                lines=6
            )

            generate_btn = gr.Button(
                "🚀 Generate",
                elem_id="generate_btn",
                variant="primary"
            )

        # RIGHT PANEL
        with gr.Column(scale=1):

            output_image = gr.Image(
                label="Generated Image",
                type="pil",
                height=500
            )

    generate_btn.click(
        fn=generate_image,
        inputs=[api_key, prompt, model_choice],
        outputs=output_image
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(), css=css)