import os
import random

import gradio as gr
import numpy as np
from dotenv import load_dotenv
from gradio_client import Client
from PIL import Image

# Load environment variables
load_dotenv()

# Client configuration
HF_TOKEN = os.getenv("HF_TOKEN")

# --- UI Constants and Helpers ---
MAX_SEED = np.iinfo(np.int32).max


def create_prompt(params):
    """
    Generate a prompt string for sending to the client

    Args:
        params (dict): Dictionary of parameters used for prompt generation
            - label (str): Comma-separated button texts, will be split into array of 4 labels
            - detail (str): Design details
            - shape (str): Button shape
            - layout (str): Layout arrangement
            - background (str): Background setting

    Returns:
        str: Complete prompt string for image generation
    """
    label_str = params.get("label", "Button")
    detail = params.get("detail", "")
    shape = params.get("shape", "rounded")
    layout = params.get("layout", "horizontal_3")
    background = params.get("background", "natural")

    # Split labels by comma, take up to number of positions
    labels = [l.strip() for l in label_str.split(",") if l.strip()]
    if labels:
        # If there are labels, repeat to up to 4
        while len(labels) < 4:
            labels.extend(labels)

    layout_config = {
        "horizontal_3": {
            "positions": ["top", "middle", "bottom"],
            "desc": "3 horizontal rows: top, middle, bottom",
        },
        "vertical_2": {
            "positions": [
                "left middle-valign small text",
                "right middle-valign small text",
            ],
            "desc": "2 vertical tall buttons arranged side by side, vertical rectangular shape, 2 columns layout, horizontal text, small text orientation is horizontal",
        },
        "box_2x2": {
            "positions": ["left-top", "right-top", "left-bottom", "right-bottom"],
            "desc": "2x2 grid: left-top, right-top, left-bottom, right-bottom",
        },
    }
    shape_descriptions = {
        "box": "box shape",
        "rounded": "rounded corners",
        "oval": "oval shape",
        "free": "organic freeform shape",
    }

    config = layout_config.get(layout, layout_config["horizontal_3"])
    positions = config["positions"]
    layout_desc = config["desc"]
    shape_desc = shape_descriptions.get(shape, "rounded corners")

    base_prompt = (
        f"Create {len(positions)} {detail} button designs in a 1024x1024 image.\n"
    )
    base_prompt += f"Arranged in {layout_desc}.\n"
    base_prompt += f"{shape_desc.capitalize()}.\n"
    if detail:
        base_prompt += f"{detail.capitalize()} aesthetic: detailed visual elements and color palette.\n"
    # Describe each button with position and label
    for i, pos in enumerate(positions):
        if i < len(labels):
            base_prompt += f'{pos} button: "{labels[i]}".\n'
        else:
            base_prompt += f"{pos} button: empty text rectangle.\n"
    base_prompt += "Each button is a different design variation exploring the theme.\n"
    if background == "natural":
        base_prompt += "Natural background with subtle textures and ambient lighting.\n"
    elif background == "white":
        base_prompt += "plain white only background.\n"
    elif background == "black":
        base_prompt += "plain black only background.\n"
    else:
        base_prompt += "Clean background with proper lighting.\n"
    base_prompt += "Ultra HD, 4K, cinematic composition"
    return base_prompt


def call_client(
    prompt, seed, randomize_seed, aspect_ratio, num_inference_steps, hf_token=None
):
    """
    Call the gradio client for image generation

    Args:
        prompt (str): Prompt text for image generation
        seed (int): Random seed value
        randomize_seed (bool): Whether to randomize the seed
        aspect_ratio (str): Image aspect ratio (e.g., "1:1", "16:9", "4:3")
        num_inference_steps (int): Number of inference steps (4-28)

    Returns:
        tuple: (image, seed, error) - Image object, used seed, error message
    """
    try:
        # Map aspect_ratio to resolution
        resolution_map = {
            "1:1": "1024x1024 ( 1:1 )",
            "16:9": "1280x720 ( 16:9 )",
            "9:16": "720x1280 ( 9:16 )",
        }
        resolution = resolution_map.get(aspect_ratio, "1024x1024 ( 1:1 )")

        client = Client("Tongyi-MAI/Z-Image-Turbo", hf_token=hf_token)
        result = client.predict(
            prompt=prompt,
            resolution=resolution,
            seed=seed,
            steps=num_inference_steps,
            shift=3.0,
            random_seed=randomize_seed,
            gallery_images=[],
            api_name="/generate",
        )
        # Assume result is PIL Image
        return result, seed, None
    except Exception as e:
        return None, seed, str(e)


# --- Main Inference Logic ---
# Define outside Blocks or define within Blocks to pass to Examples
def run_inference_engine(
    label,
    detail,
    shape,
    layout,
    background,
    seed,
    randomize_seed,
    aspect_ratio,
    guidance_scale,
    num_inference_steps,
    request: gr.Request,
):
    """
    generate UI button images

    Args:
        label (str): Text to display on the button,allow empty
        detail (str): Detailed design prompt
        shape (str): Button shape ("box", "rounded", "oval", "free")
        layout (str): Layout arrangement ("horizontal_3:3x1", "vertical_2:1x2", "box_2x2:2x2")
        background (str): Background setting ("natural", "white", "black")
        seed (int): Random seed value
        randomize_seed (bool): Whether to randomize the seed
        aspect_ratio (str): Image aspect ratio (use 1:1)
        guidance_scale (float): Guidance scale (use 1),no need to change
        num_inference_steps (int): Number of inference steps (use 8),no need to change

    Yields:
        tuple: (image, seed, status_message) - Generated image, used seed, status message
    """
    hf_token = HF_TOKEN
    if request:
        if hasattr(request, "headers"):
            if hasattr(request.headers, "authorization"):
                hf_token = request.headers.authorization
                hf_token = hf_token.replace("Bearer", "").strip()
        # print(hf_token)

    yield None, seed, "Generating..."
    prompt_params = {
        "label": label,
        "detail": detail,
        "shape": shape,
        "layout": layout,
        "background": background,
    }
    prompt = create_prompt(prompt_params)

    # Debug: Print the generated prompt
    # print(f"Generated prompt: {prompt}")

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

    image, generated_seed, error = call_client(
        prompt, seed, randomize_seed, aspect_ratio, num_inference_steps, hf_token
    )

    if image is None:
        gr.Warning(f"Error: {error}")
        yield None, generated_seed, f"Error: {error}"
    else:
        # Convert image to WebP format
        original_image = image[0][0]["image"]
        yield (original_image, generated_seed, "")


# --- UI Customization ---
css = """
body { font-family: 'Helvetica Neue', Arial, sans-serif; }
#col-container { max-width: 1200px; margin: 0 auto; padding: 20px; }
h1 { text-align: center; font-weight: 800; color: #333; margin-bottom: 0.5em; }
.subtitle { text-align: center; color: #666; margin-bottom: 2em; }

.generate-btn {
    background: linear-gradient(90deg, #6366f1 0%, #a855f7 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: bold !important;
    font-size: 1.2em !important;
    padding: 20px !important;
    border-radius: 12px !important;
    transition: all 0.3s ease;
    height: 100% !important;
    min-height: 100px;
}
.generate-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(99, 102, 241, 0.4);
}

.examples-container table {
    font-size: 0.85em !important;
    margin-bottom: 0 !important;
}
.examples-container td {
    padding: 4px 8px !important;
    white-space: nowrap;
    overflow: hidden;
    text-overflow: ellipsis;
    max-width: 150px;
}
.examples-container label {
    font-weight: bold;
    color: #555;
    margin-bottom: 5px;
    display: block;
}

#result-gallery {
    border-radius: 12px;
    box-shadow: 0 10px 30px rgba(0,0,0,0.1);
    border: 1px solid #eee;
}
.input-group {
    background: #f9fafb;
    padding: 15px;
    border-radius: 10px;
    border: 1px solid #e5e7eb;
    margin-bottom: 15px;
}
"""

theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="slate", radius_size="md")

# Examples data
example_data = [
    [
        "Start,Option,Exit",
        "Neon glowing cyberpunk, blue/purple gradient",
        "box",
        "horizontal_3",
        "black",
        "examples/start.webp",
    ],
    [
        "Buy",
        "Luxury gold texture, minimal elegant, serif",
        "rounded",
        "vertical_2",
        "white",
        "examples/buy.webp",
    ],
    [
        "RPG,R,P,G",
        "Wood texture, steel rim, fantasy game style",
        "free",
        "box_2x2",
        "natural",
        "examples/rpg.webp",
    ],
    [
        "Submit",
        "Modern flat design, blue gradient, clean minimal style",
        "rounded",
        "horizontal_3",
        "white",
        None,
    ],
]

with gr.Blocks(css=css, theme=theme, title="UI Button Generator MCP") as demo:
    gr.Markdown("# 🎨 AI UI Button Generator")
    gr.Markdown(
        "<div class='subtitle'><p>UI button material and design concept generation tool using Z-Image-Turbo(<a href='https://huggingface.co/docs/hub/spaces-zerogpu'>Zero-GPU</a>)</p><p>Web and MCP without Header-Authorization,only few time you can try zero-gpu</p></div>"
    )

    with gr.Column(elem_id="col-container"):
        # --- 1. Definition Phase (render=False) ---
        # Create components referenced by Examples first.
        # With render=False, they are not yet displayed on screen.

        # Input section
        label = gr.Textbox(
            label="Button Text",
            placeholder="Start, OK...",
            value="Start",
            info="Text inside the button",
            render=False,
        )
        detail = gr.Textbox(
            label="Detail Prompt",
            placeholder="Design details...",
            value="Pirate theme, wood texture, gold aesthetic",
            lines=4,
            info="Design details",
            render=False,
        )
        shape = gr.Dropdown(
            label="Shape",
            choices=["box", "rounded", "oval", "free"],
            value="rounded",
            info="Shape",
            render=False,
        )
        layout = gr.Radio(
            label="Layout",
            choices=["horizontal_3", "vertical_2", "box_2x2"],
            value="horizontal_3",
            info="Layout arrangement",
            render=False,
        )
        background = gr.Dropdown(
            label="Background",
            choices=["natural", "white", "black"],
            value="natural",
            info="Background",
            render=False,
        )

        # Advanced Settings section
        seed = gr.Slider(
            label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, render=False
        )
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True, render=False)
        aspect_ratio = gr.Radio(
            label="Aspect Ratio",
            choices=["1:1", "16:9", "9:16"],
            value="1:1",
            render=False,
        )
        guidance_scale = gr.Slider(
            label="Guidance Scale", minimum=1.0, maximum=5.0, value=1.0, render=False
        )
        num_inference_steps = gr.Slider(
            label="Steps", minimum=4, maximum=28, value=8, render=False
        )

        # Output section
        result = gr.Image(
            label="Output Image",
            show_label=False,
            type="pil",
            elem_id="result-image",
            height=600,
            render=False,
        )
        status_msg = gr.Markdown(render=False)

        # --- 2. Layout Construction Phase (.render()) ---

        with gr.Row(equal_height=False):
            # Left column
            with gr.Column(scale=1, min_width=400):
                with gr.Group(elem_classes="input-group"):
                    gr.Markdown("### 📝 Basic Settings")
                    label.render()
                    detail.render()

                with gr.Group(elem_classes="input-group"):
                    gr.Markdown("### 🎨 Style & Layout")
                    with gr.Row():
                        shape.render()
                        background.render()
                    layout.render()
                run_button = gr.Button(
                    "✨ Generate\nButtons",
                    variant="primary",
                    elem_classes="generate-btn",
                    scale=1,
                )

                with gr.Accordion("⚙️ Advanced Settings", open=False):
                    seed.render()
                    randomize_seed.render()
                    aspect_ratio.render()
                    with gr.Row():
                        guidance_scale.render()
                        num_inference_steps.render()

            # Right column
            with gr.Column(scale=1):
                gr.Markdown("### 🖼️ Generated Button")
                result.render()
                status_msg.render()
                # Row for button and Examples
                # gr.Markdown("**Quick Presets (Click to try)**")
                # Initialize here with fn, inputs, outputs
                gr.Examples(
                    examples=example_data,
                    fn=run_inference_engine,
                    inputs=[label, detail, shape, layout, background, result],
                    outputs=[result, seed, status_msg],
                    examples_per_page=3,
                    run_on_click=False,
                    cache_examples=False,
                )

    # --- 3. Event Binding ---
    # Examples already has fn, so only define the main button click event

    run_button.click(
        fn=run_inference_engine,
        inputs=[
            label,
            detail,
            shape,
            layout,
            background,
            seed,
            randomize_seed,
            aspect_ratio,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, seed, status_msg],
    )

if __name__ == "__main__":
    demo.launch(show_error=True, mcp_server=True, share=True)