"""
Logo Creator - Generate professional logos using AI
Self-contained app using FLUX.1-schnell with ZeroGPU on HuggingFace Spaces
"""

import os
import random

import gradio as gr
import numpy as np

MAX_SEED = np.iinfo(np.int32).max
MODEL_ID = "black-forest-labs/FLUX.1-schnell"


def generate_logo(
    prompt: str,
    style: str,
    color_scheme: str,
    background: str,
    seed: int,
    randomize_seed: bool,
    width: int,
    height: int,
    num_steps: int,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a logo based on user inputs."""
    import torch
    from diffusers import FluxPipeline

    if not prompt.strip():
        raise gr.Error("Please enter a description for your logo.")

    # Build enhanced prompt
    parts = [
        f"A professional logo design: {prompt.strip()}",
        f"{style} style" if style != "Auto" else "",
        f"{color_scheme} color scheme" if color_scheme != "Auto" else "",
        f"{background} background" if background != "Auto" else "",
        "clean vector art, high quality, centered composition, professional logo design",
    ]
    full_prompt = ", ".join(p for p in parts if p)

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    pipe = FluxPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.bfloat16,
    )
    pipe.to("cuda")

    generator = torch.Generator("cuda").manual_seed(seed)

    image = pipe(
        prompt=full_prompt,
        width=width,
        height=height,
        num_inference_steps=num_steps,
        generator=generator,
        guidance_scale=0.0,
    ).images[0]

    return image, seed


# Apply ZeroGPU decorator if available
try:
    import spaces
    generate_logo = spaces.GPU(duration=120)(generate_logo)
    print("ZeroGPU enabled")
except ImportError:
    print("Running without ZeroGPU (local mode)")


# --- UI ---
STYLE_CHOICES = [
    "Auto", "Minimalist", "Modern", "Vintage", "Geometric", "Abstract",
    "Flat", "3D", "Hand-drawn", "Corporate", "Playful", "Elegant", "Tech", "Nature",
]

COLOR_CHOICES = [
    "Auto", "Monochrome", "Blue tones", "Red tones", "Green tones",
    "Gold and black", "Pastel", "Vibrant", "Earth tones", "Gradient", "Neon",
]

BG_CHOICES = ["Auto", "White", "Black", "Transparent-style", "Gradient"]

EXAMPLES = [
    ["A mountain peak with a rising sun for a travel company called Horizon", "Minimalist", "Blue tones", "White", 42, True, 1024, 1024, 4],
    ["A fox head for a tech startup called FoxByte", "Geometric", "Auto", "White", 42, True, 1024, 1024, 4],
    ["Coffee cup with steam forming a book for ReadBrew", "Modern", "Earth tones", "Auto", 42, True, 1024, 1024, 4],
    ["A rocket launching from a laptop for a SaaS company", "Flat", "Vibrant", "White", 42, True, 1024, 1024, 4],
    ["An elegant letter M with floral accents for a luxury brand", "Elegant", "Gold and black", "Black", 42, True, 1024, 1024, 4],
]

with gr.Blocks(
    theme=gr.themes.Soft(),
    title="Logo Creator",
    css=".main-header { text-align: center; margin-bottom: 0.5em; } .generate-btn { min-height: 50px !important; }",
) as demo:
    gr.Markdown("# Logo Creator", elem_classes="main-header")
    gr.Markdown("Generate professional logos with AI. Describe your logo idea and customize the style.")

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Logo Description",
                placeholder="Describe the logo you want to create...",
                lines=3,
            )
            style = gr.Dropdown(choices=STYLE_CHOICES, value="Auto", label="Style")
            color_scheme = gr.Dropdown(choices=COLOR_CHOICES, value="Auto", label="Color Scheme")
            background = gr.Dropdown(choices=BG_CHOICES, value="Auto", label="Background")

            with gr.Accordion("Advanced Settings", open=False):
                seed = gr.Slider(minimum=0, maximum=MAX_SEED, step=1, value=42, label="Seed")
                randomize_seed = gr.Checkbox(value=True, label="Randomize Seed")
                with gr.Row():
                    width = gr.Slider(minimum=256, maximum=1024, step=64, value=1024, label="Width")
                    height = gr.Slider(minimum=256, maximum=1024, step=64, value=1024, label="Height")
                num_steps = gr.Slider(minimum=1, maximum=8, step=1, value=4, label="Inference Steps")

            generate_btn = gr.Button("Generate Logo", variant="primary", elem_classes="generate-btn")

        with gr.Column(scale=1):
            output_image = gr.Image(label="Generated Logo", type="pil")
            output_seed = gr.Number(label="Seed Used", interactive=False)

    generate_btn.click(
        fn=generate_logo,
        inputs=[prompt, style, color_scheme, background, seed, randomize_seed, width, height, num_steps],
        outputs=[output_image, output_seed],
    )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[prompt, style, color_scheme, background, seed, randomize_seed, width, height, num_steps],
        outputs=[output_image, output_seed],
        fn=generate_logo,
        cache_examples=False,
    )

    gr.Markdown("---\n*Powered by [FLUX.1-schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell) with ZeroGPU*")

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