"""
Dispatch AI — Wallpaper Engine
Gradio app generating phone wallpapers with FLUX.1-schnell.
Presets: UAE Sunset, Cyberpunk, Minimal Dark, Arabic Calligraphy + more.
"""

import os
import io
import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image, ImageDraw

# --- Configuration -----------------------------------------------------------
HF_TOKEN = os.environ.get("HF_TOKEN", None)
MODEL_ID = "black-forest-labs/FLUX.1-schnell"

client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)

BG_COLOR = "#0A0F1A"
ACCENT = "#1FE0E6"

# Aspect ratio -> (width, height)
ASPECT_RATIOS = {
    "9:16 (Phone Wallpaper)": (720, 1280),
    "9:16 HD (Phone)": (1080, 1920),
    "1:1 (Square)": (1024, 1024),
    "16:9 (Wide)": (1280, 720),
    "4:5 (Portrait)": (1024, 1280),
}

# Presets: name -> (prompt, emoji)
PRESETS = {
    "UAE Sunset": (
        "A breathtaking UAE desert sunset, golden sand dunes rolling under a vibrant orange and purple sky, "
        "silhouette of a lone camel walking, distant wind tower, cinematic lighting, ultra detailed, 4k wallpaper",
        "🏜️",
    ),
    "Cyberpunk City": (
        "A neon-soaked cyberpunk city at night, towering skyscrapers with glowing cyan and magenta lights, "
        "flying cars, rain-slicked streets reflecting neon, blade runner aesthetic, ultra detailed wallpaper",
        "🌃",
    ),
    "Minimal Dark": (
        "Minimalist dark abstract wallpaper, deep black background with subtle geometric cyan lines, "
        "smooth gradients, modern elegant design, 4k",
        "⬛",
    ),
    "Arabic Calligraphy": (
        "Beautiful Arabic calligraphy art wallpaper, elegant flowing script in gold on dark textured background, "
        "intricate decorative borders, luxurious design, 4k",
        "✒️",
    ),
    "Desert Stars": (
        "A serene Arabian desert at night under a sky full of stars, Milky Way galaxy visible, "
        "soft sand dunes in foreground, peaceful cosmic atmosphere, ultra detailed wallpaper",
        "🌌",
    ),
    "Dubai Skyline": (
        "Dubai skyline at golden hour, Burj Khalifa towering above modern skyscrapers, warm sunlight, "
        "palm trees, luxury cars on Sheikh Zayed Road, cinematic, ultra detailed, 4k wallpaper",
        "🏙️",
    ),
    "Ocean Waves": (
        "Tranquil ocean waves wallpaper, deep blue water with white foam, sunlight sparkling on the surface, "
        "minimalist calming aesthetic, top-down view, ultra detailed, 4k",
        "🌊",
    ),
    "Mountain Peak": (
        "Majestic snow-capped mountain peak at dawn, alpine glow, dramatic clouds, pristine landscape, "
        "ultra detailed nature wallpaper, 4k",
        "🏔️",
    ),
    "Abstract Neon": (
        "Abstract neon wallpaper, flowing liquid metal shapes in cyan and purple, dark background, "
        "futuristic 3D render, ultra sharp, 4k",
        "🎨",
    ),
    "Falcon Majesty": (
        "A majestic falcon perched on a rock in the Arabian desert, piercing eyes, detailed feathers, "
        "warm sunset light, ultra detailed wildlife wallpaper, 4k",
        "🦅",
    ),
}


def generate_wallpaper(prompt, aspect, preset, num_steps):
    """Generate a wallpaper via FLUX.1-schnell."""
    if preset and preset in PRESETS:
        final_prompt = PRESETS[preset][0]
    elif prompt and prompt.strip():
        final_prompt = prompt.strip()
    else:
        final_prompt = PRESETS["UAE Sunset"][0]

    w, h = ASPECT_RATIOS.get(aspect, (720, 1280))
    steps = int(num_steps) if num_steps else 4

    try:
        image = client.text_to_image(
            final_prompt,
            width=w,
            height=h,
            num_inference_steps=steps,
        )
        if not isinstance(image, Image.Image):
            image = Image.open(io.BytesIO(image)) if hasattr(image, "read") else Image.open(image)
        return image, "✅ Wallpaper generated successfully!"
    except Exception as e:
        img = Image.new("RGB", (w, h), BG_COLOR)
        d = ImageDraw.Draw(img)
        d.text((w // 4, h // 2), f"Error: {str(e)[:60]}", fill=ACCENT)
        return img, f"❌ Error: {str(e)}"


def download_image(img):
    if img is None:
        return None
    path = os.path.join(os.getcwd(), "wallpaper_output.png")
    img.save(path)
    return path


def pick_preset_prompt(preset_name):
    return PRESETS.get(preset_name, ("", ""))[0]


# --- UI -----------------------------------------------------------------------
CSS = """
#dispatch-header h1 {
    color: #FFFFFF; font-size: 2.2rem; margin: 0;
    background: linear-gradient(90deg, #1FE0E6 0%, #FFFFFF 60%);
    -webkit-background-clip: text; -webkit-text-fill-color: transparent;
}
#dispatch-header p { color: #1FE0E6; font-size: 1.05rem; margin: 6px 0 0 0; }
.dispatch-footer { text-align: center; color: #8A8F9C; font-size: 0.9rem; padding-top: 8px; }
"""

with gr.Blocks(
    title="Dispatch AI — Wallpaper Engine",
    theme=gr.themes.Base(
        primary_hue="cyan",
        secondary_hue="cyan",
        neutral_hue="slate",
        font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui"],
    ).set(
        body_background_fill="#0A0F1A",
        body_background_fill_dark="#0A0F1A",
        body_text_color="#FFFFFF",
        body_text_color_dark="#FFFFFF",
        block_background_fill="#0E1424",
        block_background_fill_dark="#0E1424",
        block_border_color="#1FE0E6",
        block_border_width="1px",
        block_label_text_color="#1FE0E6",
        block_title_text_color="#1FE0E6",
        button_primary_background_fill="#1FE0E6",
        button_primary_background_fill_dark="#1FE0E6",
        button_primary_text_color="#0A0F1A",
        button_primary_border_color="#1FE0E6",
        input_background_fill="#0E1424",
        input_background_fill_dark="#0E1424",
        input_border_color="#1FE0E6",
        input_border_width="1px",
    ),
    css=CSS,
) as demo:
    with gr.Column(elem_id="dispatch-header"):
        gr.Markdown(
            """
            # Dispatch AI — Wallpaper Engine
            Generate stunning phone wallpapers with FLUX.1-schnell · Dispatch AI (FZE) · UAE
            """
        )

    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="Describe your wallpaper... or pick a preset below",
                lines=3,
            )
            aspect_select = gr.Radio(
                list(ASPECT_RATIOS.keys()),
                label="Aspect Ratio",
                value="9:16 (Phone Wallpaper)",
            )
            preset_select = gr.Dropdown(
                list(PRESETS.keys()),
                label="Preset",
                value=None,
                info="Select a preset to auto-fill the prompt",
            )
            steps_slider = gr.Slider(
                minimum=1, maximum=8, value=4, step=1,
                label="Inference Steps (schnell works well at 4)",
            )
            with gr.Row():
                load_preset_btn = gr.Button("Load Preset", variant="secondary")
                generate_btn = gr.Button("🎨 Generate Wallpaper", variant="primary")
        with gr.Column(scale=2):
            output_image = gr.Image(
                label="Generated Wallpaper",
                type="pil",
                show_download_button=True,
            )
            status_box = gr.Textbox(label="Status", interactive=False)
            download_btn = gr.Button("⬇️ Download PNG", variant="secondary")
            download_file = gr.File(label="Download")

    # Events
    load_preset_btn.click(pick_preset_prompt, inputs=preset_select, outputs=prompt_input)
    generate_btn.click(
        generate_wallpaper,
        inputs=[prompt_input, aspect_select, preset_select, steps_slider],
        outputs=[output_image, status_box],
    )
    download_btn.click(download_image, inputs=output_image, outputs=download_file)

    gr.Examples(
        examples=[[v[0], "9:16 (Phone Wallpaper)", k, 4] for k, v in PRESETS.items()],
        inputs=[prompt_input, aspect_select, preset_select, steps_slider],
        label="Preset Examples — click to load",
    )

    gr.Markdown(
        """
        <div class="dispatch-footer">
        © 2026 Dispatch AI (FZE) · UAE · License 10818 · Model: FLUX.1-schnell by Black Forest Labs
        </div>
        """
    )

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