app.py

"""
Text-to-Game Generator
Pipeline:
  Theme --> [Groq Llama] --> HTML5 game code (with sprite_NAME.png refs)
  Theme --> [Groq Llama acting as Z-Image-Engineer] --> cinematic image prompts
         --> [FLUX.1-dev via fal-ai] --> sprite images
         --> injected as base64 into game HTML

Secrets needed:
  GROQ_API_KEY  - console.groq.com (free, no credit card)
  HF_TOKEN      - huggingface.co/settings/tokens (for FLUX.1-dev via fal-ai)
"""

import os
import re
import io
import base64
import traceback
import time
import gradio as gr
from openai import OpenAI
import requests
from urllib.parse import quote
from huggingface_hub import InferenceClient
from PIL import Image

# ---------------------------------------------------------------------------
# Clients
# ---------------------------------------------------------------------------

GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
HF_TOKEN     = os.environ.get("HF_TOKEN", "")

CODE_MODEL   = "llama-3.1-8b-instant"          # Groq — game code
PROMPT_MODEL = "llama-3.3-70b-versatile"       # Groq — image prompt enhancement
IMAGE_MODEL  = "black-forest-labs/FLUX.1-dev"  # fal-ai — image generation
POLLINATIONS_URL = "https://image.pollinations.ai/prompt/{prompt}?width={w}&height={h}&model=flux&nologo=true&seed={seed}"

# Canvas dimensions — iframe is fixed to these so game fits perfectly
CANVAS_W = 800
CANVAS_H = 450


def get_groq_client():
    if not GROQ_API_KEY:
        raise ValueError(
            "GROQ_API_KEY not set. "
            "Get a free key at console.groq.com and add it as a Space secret."
        )
    return OpenAI(
        base_url="https://api.groq.com/openai/v1",
        api_key=GROQ_API_KEY,
    )


# ---------------------------------------------------------------------------
# Z-Image-Engineer system prompt (from BennyDaBall/Qwen3-4b-Z-Image-Engineer-V4)
# ---------------------------------------------------------------------------

Z_ENGINEER_SYSTEM = (
    "Interpret the user seed as production intent, then build a definitive 200-250 word "
    "single-paragraph image prompt that preserves every explicit constraint while intelligently "
    "expanding missing details. First infer the core subject, action, setting, and emotional tone; "
    "treat these as non-negotiable anchors. Then enhance with precise visual staging "
    "(explicit foreground, midground, background), clear visual hierarchy and eye path, "
    "physically plausible lighting (source, direction, softness, color temperature), and optical "
    "strategy (if lens/aperture are provided, preserve exactly; if absent, choose fitting lens and "
    "aperture and imply their depth-of-field effect). Integrate organic, manufactured, and "
    "environmental textures with realistic material behavior, add motion/atmospheric cues only "
    "when they support the scene, and apply a coherent color grade consistent with mood and "
    "environment. Output ONLY the image prompt paragraph. No explanation, no preamble."
)

# ---------------------------------------------------------------------------
# Game type configs — Platformer and Top-Down Shooter only
# ---------------------------------------------------------------------------

GAME_TYPES = {
    "Platformer": {
        "description": "Reach the goal platform while avoiding wandering monsters.",
        "prompt_template": (
            "Create a complete, self-contained HTML5 platformer game with the theme: {theme}.\n"
            "EXACT GAME RULES — implement every one precisely:\n"
            "- Canvas: id='gameCanvas', size 800x450.\n"
            "- PLATFORMS: at least 6 platforms drawn with ctx.drawImage(platformImg, x, y, w, h). "
            "  Include one full-width ground platform at y=420 height=20. "
            "  Place 5 elevated platforms at varied x/y positions. "
            "  Use new Image() with src='sprite_platform.png' for ALL platforms.\n"
            "- TARGET/GOAL: a glowing star or chest sprite at the top-right platform. "
            "  Use new Image() with src='sprite_goal.png'. "
            "  When player bounding box overlaps goal: show WIN screen with score and Restart button.\n"
            "- PLAYER: starts bottom-left. Size 40x40. "
            "  Move left/right with A/D or ArrowLeft/ArrowRight (speed 4). "
            "  Jump with W, ArrowUp, or Space (velY = -12, only when grounded). "
            "  Gravity: velY += 0.5 every frame. "
            "  Platform collision: set grounded=false BEFORE loop; inside loop if player feet hit platform top set grounded=true velY=0. "
            "  Keep player inside canvas horizontally.\n"
            "- MONSTERS: 3 monsters, each patrolling back-and-forth on its own platform. "
            "  Size 32x32. Speed 1.5 px/frame. Reverse direction at platform edges. "
            "  Use new Image() with src='sprite_enemy.png'. "
            "  If monster bounding box overlaps player: player lives -= 1, respawn player at start. "
            "  0 lives = GAME OVER screen with Restart button.\n"
            "- HUD: lives top-left, score top-right.\n"
            "- IMAGES: declare all at top of script before anything else: "
            "  const playerImg = new Image(); playerImg.src = 'sprite_player.png'; "
            "  const bgImg = new Image(); bgImg.src = 'sprite_background.png'; "
            "  const enemyImg = new Image(); enemyImg.src = 'sprite_enemy.png'; "
            "  const platformImg = new Image(); platformImg.src = 'sprite_platform.png'; "
            "  const goalImg = new Image(); goalImg.src = 'sprite_goal.png'; "
            "  Use Promise.all([loadImg(playerImg),loadImg(bgImg),loadImg(enemyImg),loadImg(platformImg),loadImg(goalImg)]).then(startGame).\n"
            "- Define ALL functions at TOP LEVEL, not inside Promise.then or startGame.\n"
            "- Add window keydown/keyup listeners to a keys Set inside startGame().\n"
            "- NO external libraries, NO CDN links.\n"
            "Output ONLY the raw HTML. No explanation, no markdown fences."
        ),
    },
    "Top-Down Shooter": {
        "description": "Shoot monsters coming from the top before they reach you.",
        "prompt_template": (
            "Create a complete, self-contained HTML5 top-down shooter game with the theme: {theme}.\n"
            "EXACT GAME RULES — implement every one precisely:\n"
            "- Canvas: id='gameCanvas', size 800x450.\n"
            "- MONSTERS spawn at random x positions along the TOP edge (y=0) and move DOWNWARD only (y += speed each frame).\n"
            "- Monster size: 32x32. Monster speed: 1-2 px/frame (slower than player speed of 4).\n"
            "- PLAYER starts at bottom-center, moves left/right only with A/D or ArrowLeft/ArrowRight keys.\n"
            "- Player size: 48x48. Player speed: 4 px/frame. Keep player inside canvas bounds.\n"
            "- LEFT MOUSE CLICK fires one bullet from player position toward the click point.\n"
            "  Bullet speed: 10 px/frame. Remove bullets that leave canvas.\n"
            "- BULLET HIT: if a bullet rect overlaps a monster rect, remove BOTH the bullet and the monster. Score += 10.\n"
            "- MONSTER COLLISION: if a monster rect overlaps the player rect, remove the monster. Player health -= 1.\n"
            "- MONSTER ESCAPED: if a monster reaches y > canvas.height remove it (no health loss).\n"
            "- New monsters spawn every 90 frames. Spawn rate increases every 500 points.\n"
            "- HUD: draw score top-left, health top-right (show as hearts or number).\n"
            "- GAME OVER when health reaches 0: show score and a Restart button.\n"
            "- NO gravity, NO velY += 0.5, NO grounded, NO jumping.\n"
            "- Use new Image() with src='sprite_player.png' for the player (48x48).\n"
            "- Use new Image() with src='sprite_background.png' for the background (full canvas).\n"
            "- Use new Image() with src='sprite_enemy.png' for monsters (32x32).\n"
            "- Declare all images at top of script, use Promise.all to wait before starting gameLoop.\n"
            "- Define all functions at TOP LEVEL, not inside Promise.then or startGame.\n"
            "- Add window keydown/keyup listeners to a keys Set inside startGame().\n"
            "- NO external libraries, NO CDN links.\n"
            "Output ONLY the raw HTML. No explanation, no markdown fences."
        ),
    },
}

GAME_TYPE_NAMES = list(GAME_TYPES.keys())

THEME_EXAMPLES = {
    "Platformer":       [["Jungle temple with ancient traps"], ["Neon cyberpunk rooftops"], ["Underwater pirate shipwreck"]],
    "Top-Down Shooter": [["Ancient Egyptian tomb raid with cursed mummies"], ["Alien desert invasion"], ["Viking village under siege"]],
}

# ---------------------------------------------------------------------------
# Step 1: Generate image prompts via Z-Image-Engineer (Groq)
# ---------------------------------------------------------------------------

def generate_image_prompts(theme: str, game_type: str) -> dict:
    client = get_groq_client()
    if game_type == "Top-Down Shooter":
        bg_style   = "bird's eye overhead view, top-down 2D game background, viewed from directly above, like a map"
        char_style = (
            "top-down overhead game sprite, viewed from directly above, "
            "character body seen from above like looking straight down, "
            "head at top shoulders below, like Metal Slug or GTA 2 sprite angle, "
            "pure black background, character only"
        )
    else:
        bg_style   = "2D side-scrolling platformer background, horizontal landscape viewed from the side"
        char_style = (
            "2D side-view platformer game sprite, character facing right, "
            "full body visible from the side like Super Mario or Mega Man, "
            "classic side-scrolling game art style, "
            "pure black background, character only"
        )

    seeds = {
        "sprite_player.png": (
            f"{char_style}, {theme} theme, "
            f"vibrant colors, strong clear silhouette, 64x64 pixel style, "
            f"single character centered, no scenery no ground no environment"
        ),
        "sprite_background.png": (
            f"{bg_style}, {theme} theme, "
            f"wide atmospheric scene, game art style, 800x450, "
            f"no characters no sprites, environment only"
        ),
        "sprite_enemy.png": (
            f"{char_style}, {theme} theme, enemy monster villain, "
            f"menacing threatening design, strong clear silhouette, 64x64 pixel style, "
            f"single enemy centered, no scenery no ground no environment"
        ),
    }
    if game_type == "Platformer":
        seeds["sprite_platform.png"] = (
            f"2D side-view pixel-art platform tile, {theme} theme, "
            f"rectangular solid surface like a game platform, stone or wood texture, "
            f"pure black background, platform only, 128x24 pixel style"
        )
        seeds["sprite_goal.png"] = (
            f"2D side-view pixel-art treasure chest or glowing star goal, {theme} theme, "
            f"glowing clearly visible reward item, "
            f"pure black background, item only, 40x40 pixel style"
        )
    prompts = {}
    for sprite_name, seed in seeds.items():
        try:
            response = client.chat.completions.create(
                model=PROMPT_MODEL,
                messages=[
                    {"role": "system", "content": Z_ENGINEER_SYSTEM},
                    {"role": "user",   "content": seed},
                ],
                max_tokens=400,
                temperature=0.8,
            )
            prompts[sprite_name] = response.choices[0].message.content.strip()
        except Exception as exc:
            print(f"[Z-Engineer] Failed {sprite_name}: {exc}")
            prompts[sprite_name] = seed
    return prompts

# ---------------------------------------------------------------------------
# Step 2: Generate images via FLUX.1-dev (fal-ai)
# ---------------------------------------------------------------------------

def _remove_background(pil_image: Image.Image) -> Image.Image:
    """Remove background from sprite using rembg library."""
    try:
        from rembg import remove
        buf_in = io.BytesIO()
        pil_image.save(buf_in, format="PNG")
        buf_in.seek(0)
        buf_out = io.BytesIO()
        remove(buf_in.read(), output=buf_out)
        buf_out.seek(0)
        return Image.open(buf_out).convert("RGBA")
    except Exception as exc:
        print(f"[rembg] Background removal failed: {exc}")
        return pil_image


def _pil_to_data_uri(pil_image: Image.Image, size: tuple = None) -> str:
    img = pil_image.resize(size, Image.LANCZOS) if size else pil_image
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()


def _colored_placeholder(name: str) -> str:
    colours = ["#e74c3c", "#3498db", "#2ecc71", "#f39c12", "#9b59b6", "#1abc9c"]
    colour  = colours[abs(hash(name)) % len(colours)]
    label   = name.replace("sprite_", "")[:6]
    svg = (
        '<svg xmlns="http://www.w3.org/2000/svg" width="64" height="64">'
        '<rect width="64" height="64" fill="' + colour + '" rx="8"/>'
        '<text x="32" y="38" font-size="10" text-anchor="middle" fill="white">' + label + '</text>'
        '</svg>'
    )
    return "data:image/svg+xml;base64," + base64.b64encode(svg.encode()).decode()


def _generate_via_hf(prompt: str) -> Image.Image:
    """Try FLUX.1-dev via HF auto-provider selection.
    HF automatically picks the fastest available provider — if one is
    quota-depleted it routes to the next available one automatically."""
    client = InferenceClient(api_key=HF_TOKEN)
    return client.text_to_image(prompt, model=IMAGE_MODEL)


def _generate_via_pollinations(prompt: str, sprite_name: str, is_bg: bool) -> Image.Image:
    """Fallback: Pollinations.AI, free, no key needed."""
    w, h = (CANVAS_W, CANVAS_H) if is_bg else (64, 64)
    seed = abs(hash(sprite_name)) % 99999
    url  = POLLINATIONS_URL.format(prompt=quote(prompt), w=w, h=h, seed=seed)
    for attempt in range(3):
        try:
            resp = requests.get(url, timeout=120)
            resp.raise_for_status()
            return Image.open(io.BytesIO(resp.content))
        except Exception as exc:
            if attempt < 2:
                time.sleep(20)
            else:
                raise exc


def generate_sprites(image_prompts: dict) -> tuple:
    sprite_map = {}
    errors     = []

    for i, (sprite_name, prompt) in enumerate(image_prompts.items()):
        try:
            is_bg = "background" in sprite_name

            # Try FLUX.1-dev via fal-ai if HF_TOKEN is set
            if HF_TOKEN:
                try:
                    pil_img = _generate_via_hf(prompt)
                    provider = "FLUX.1-dev/auto"
                except Exception as fal_exc:
                    fal_err = str(fal_exc)
                    print(f"[fal-ai] Failed {sprite_name}: {fal_err} — falling back to Pollinations")
                    errors.append(f"{sprite_name} (fal-ai failed, used Pollinations): {fal_err[:60]}")
                    if i > 0:
                        time.sleep(20)
                    pil_img  = _generate_via_pollinations(prompt, sprite_name, is_bg)
                    provider = "Pollinations"
            else:
                # No HF_TOKEN — go straight to Pollinations
                if i > 0:
                    time.sleep(20)
                pil_img  = _generate_via_pollinations(prompt, sprite_name, is_bg)
                provider = "Pollinations"

            if not is_bg:
                pil_img = _remove_background(pil_img)
            size = None if is_bg else (64, 64)
            sprite_map[sprite_name] = _pil_to_data_uri(pil_img, size=size)
            print(f"[{provider}] OK: {sprite_name}")

        except Exception as exc:
            error_msg = str(exc)
            errors.append(f"{sprite_name}: {error_msg[:80]}")
            print(f"[Image] FAILED {sprite_name}: {error_msg}")
            sprite_map[sprite_name] = _colored_placeholder(sprite_name.replace(".png", ""))

    return sprite_map, errors

# ---------------------------------------------------------------------------
# Step 3: Inject sprites into HTML
# ---------------------------------------------------------------------------

def _inject_sprites(html_code: str, sprite_map: dict) -> str:
    for fname, data_uri in sprite_map.items():
        html_code = html_code.replace('"' + fname + '"', '"' + data_uri + '"')
        html_code = html_code.replace("'" + fname + "'", "'" + data_uri + "'")
    return html_code

# ---------------------------------------------------------------------------
# Step 4: Generate game code via Groq Llama
# ---------------------------------------------------------------------------

CODE_SYSTEM = (
    "You are an expert HTML5 game developer. "
    "Write a complete, working, single-file HTML5 game using canvas and vanilla JavaScript. "
    "CRITICAL RULES - copy these patterns EXACTLY as written, do not deviate: "
    "1. The VERY FIRST lines of the script must be exactly: "
    "   const canvas = document.getElementById('gameCanvas'); "
    "   const ctx = canvas.getContext('2d'); "
    "   const playerImg = new Image(); playerImg.src = 'sprite_player.png'; "
    "   const bgImg = new Image();     bgImg.src = 'sprite_background.png'; "
    "   const enemyImg = new Image();  enemyImg.src = 'sprite_enemy.png'; "
    "   const keys = new Set(); "
    "   const bullets = []; "
    "   const enemies = []; "
    "   let score = 0; let health = 3; let frameCount = 0; let gameOver = false; "
    "2. FOR TOP-DOWN SHOOTER declare player at CENTER - copy EXACTLY: "
    "   let player = {x: canvas.width/2-24, y: canvas.height/2-24, w:48, h:48, speed:4}; "
    "   DO NOT use canvas.height - 24 for player y. Player starts in CENTER not bottom. "
    "3. Use Promise.all AFTER all declarations: "
    "   function loadImg(img) { return new Promise(r => { img.onload = r; }); } "
    "   Promise.all([loadImg(playerImg), loadImg(bgImg), loadImg(enemyImg)]).then(startGame); "
    "4. startGame() adds ONLY keyboard listeners and calls gameLoop - NO click listener here: "
    "   function startGame() { "
    "     window.addEventListener('keydown', e => keys.add(e.key)); "
    "     window.addEventListener('keyup', e => keys.delete(e.key)); "
    "     canvas.addEventListener('click', onShoot); "
    "     requestAnimationFrame(gameLoop); } "
    "5. gameLoop() NEVER redeclares canvas or ctx. Draw background FIRST with FULL SIZE: "
    "   ctx.drawImage(bgImg, 0, 0, canvas.width, canvas.height); "
    "   This is critical - always pass canvas.width and canvas.height as 3rd and 4th arguments. "
    "6. FOR TOP-DOWN SHOOTER movement - 4 directions, clamp inside canvas: "
    "   if (keys.has('w')||keys.has('W')||keys.has('ArrowUp'))    player.y -= player.speed; "
    "   if (keys.has('s')||keys.has('S')||keys.has('ArrowDown'))  player.y += player.speed; "
    "   if (keys.has('a')||keys.has('A')||keys.has('ArrowLeft'))  player.x -= player.speed; "
    "   if (keys.has('d')||keys.has('D')||keys.has('ArrowRight')) player.x += player.speed; "
    "   player.x = Math.max(0, Math.min(canvas.width-player.w, player.x)); "
    "   player.y = Math.max(0, Math.min(canvas.height-player.h, player.y)); "
    "7. Bullets fire on MOUSE CLICK straight UP only - no vx component: "
    "   function onShoot(e) { if (gameOver) return; "
    "     bullets.push({x:player.x+player.w/2-4, y:player.y-16, w:8, h:16, vy:-10}); } "
    "   Each frame: bullets[i].y += bullets[i].vy; draw yellow rect; remove when y+h<0. "
    "8. Enemies spawn every 120 frames, fall straight down, speed capped at 3.5: "
    "   frameCount++; "
    "   if (frameCount % 120 === 0) enemies.push({x:Math.random()*(canvas.width-32), y:0, w:32, h:32, speed: Math.min(1+score/500, 3.5)}); "
    "   Each frame: e.y += e.speed; ctx.drawImage(enemyImg,e.x,e.y,32,32); "
    "   Remove if e.y > canvas.height. If overlaps player: health--; remove enemy. "
    "9. Draw player: ctx.drawImage(playerImg, player.x, player.y, player.w, player.h). "
    "10. GAME OVER when health<=0 - set gameOver=true then draw overlay INSIDE gameLoop: "
    "    ctx.fillStyle='rgba(0,0,0,0.7)'; ctx.fillRect(0,0,canvas.width,canvas.height); "
    "    draw GAME OVER text and score at center. "
    "    draw a green restart button rect at center+60px. "
    "    add ONE-TIME click listener for restart ONLY when gameOver becomes true: "
    "    canvas.removeEventListener('click', onShoot); "
    "    canvas.addEventListener('click', restartHandler); "
    "    then call return to stop the loop. "
    "11. restartHandler resets ALL variables and restores onShoot listener: "
    "    function restartHandler() { "
    "      canvas.removeEventListener('click', restartHandler); "
    "      player.x=canvas.width/2-24; player.y=canvas.height/2-24; "
    "      bullets.length=0; enemies.length=0; "
    "      score=0; health=3; frameCount=0; gameOver=false; "
    "      canvas.addEventListener('click', onShoot); "
    "      requestAnimationFrame(gameLoop); } "
    "12. FOR PLATFORMER: gravity velY+=0.5, jump ArrowUp/W/Space velY=-12 when grounded. "
    "    Set grounded=false BEFORE platform loop. Set true and velY=0 only on landing. "
    "    Full-width ground at y=420. "
    "Output ONLY the raw HTML - no markdown fences, no explanation."
)

def generate_game_code(game_type: str, theme: str, temperature: float, max_new_tokens: int):
    if not theme.strip():
        return "", "", "Please enter a theme first.", _placeholder_html("Enter a theme and generate a game.")

    try:
        client      = get_groq_client()
        user_prompt = GAME_TYPES[game_type]["prompt_template"].format(theme=theme.strip())

        # Code generation
        code_resp = client.chat.completions.create(
            model=CODE_MODEL,
            messages=[
                {"role": "system", "content": CODE_SYSTEM},
                {"role": "user",   "content": user_prompt},
            ],
            max_tokens=int(max_new_tokens),
            temperature=float(temperature),
        )
        code = code_resp.choices[0].message.content.strip()
        code = re.sub(r"^```[a-zA-Z]*\n?", "", code).strip()
        code = re.sub(r"\n?```$",           "", code).strip()
        if "<html" not in code.lower() and "<!doctype" not in code.lower():
            code = _wrap_in_html(code, theme)

        # Image prompts
        image_prompts = generate_image_prompts(theme.strip(), game_type)

        # Sprites
        sprite_map, sprite_errors = generate_sprites(image_prompts)

        # Inject
        final_code = _inject_sprites(code, sprite_map)

        n_real     = sum(1 for v in sprite_map.values() if "image/png" in v)
        n_fallback = len(sprite_map) - n_real

        if sprite_errors:
            status = f"Code done. Images: {n_real} FLUX, {n_fallback} fallback. Errors: {' | '.join(sprite_errors)}"
        else:
            status = f"Done! {n_real} sprite(s) by FLUX.1-dev. Click Launch Game to play."

        prompt_summary = "\n\n".join(f"**{k}:**\n{v}" for k, v in image_prompts.items())
        return final_code, prompt_summary, status, _build_preview(final_code)

    except Exception as exc:
        traceback.print_exc()
        err = str(exc)
        if "401" in err or "api_key" in err.lower():
            err = "Invalid GROQ_API_KEY. Check your key at console.groq.com."
        elif "429" in err or "rate" in err.lower():
            err = "Rate limited by Groq - wait a few seconds and try again."
        else:
            err = "Error: " + str(exc)
        return "", "", err, _placeholder_html(err)


# ---------------------------------------------------------------------------
# HTML helpers
# ---------------------------------------------------------------------------

def _placeholder_html(message: str) -> str:
    safe = message.replace("<", "&lt;").replace(">", "&gt;")
    return (
        f'<div style="display:flex;align-items:center;justify-content:center;'
        f'width:{CANVAS_W}px;height:{CANVAS_H}px;background:#0d0d0d;border-radius:12px;'
        f'border:2px dashed #333;color:#555;font-family:monospace;font-size:14px;'
        f'text-align:center;padding:24px;box-sizing:border-box;">'
        f'<pre style="margin:0;white-space:pre-wrap;">{safe}</pre></div>'
    )


def _wrap_in_html(snippet: str, theme: str) -> str:
    return (
        "<!DOCTYPE html>\n<html lang='en'>\n<head>\n"
        "<meta charset='UTF-8'><title>" + theme + "</title>\n"
        "<style>body{margin:0;background:#111;display:flex;justify-content:center;"
        "align-items:center;height:100vh;}canvas{display:block;}</style>\n"
        "</head>\n<body>\n" + snippet + "\n</body>\n</html>"
    )


def _build_preview(html_code: str) -> str:
    encoded = base64.b64encode(html_code.encode("utf-8")).decode("ascii")
    return (
        f'<div style="width:{CANVAS_W}px;height:{CANVAS_H}px;overflow:hidden;border-radius:12px;">'  
        f'<iframe src="data:text/html;base64,{encoded}" '
        f'style="width:{CANVAS_W}px;height:{CANVAS_H}px;border:none;background:#000;display:block;" '
        f'scrolling="no" '
        'sandbox="allow-scripts" title="Game Preview"></iframe></div>'
    )


def launch_game(code: str) -> str:
    if not code or not code.strip():
        return _placeholder_html("No game code yet - generate a game first.")
    return _build_preview(code)


# ---------------------------------------------------------------------------
# UI helpers
# ---------------------------------------------------------------------------

def update_type_description(game_type: str) -> str:
    return "_" + GAME_TYPES[game_type]["description"] + "_"


def get_first_theme(game_type: str) -> str:
    return THEME_EXAMPLES[game_type][0][0]


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

def build_ui():
    with gr.Blocks(title="Game Generator") as demo:

        gr.Markdown(
            "# Game Generator\n"
            "Type a theme — the AI writes the game code, generates cinematic image prompts "
            "using **Z-Image-Engineer V4** style, then **FLUX.1-dev** renders the sprites.\n\n"
            "> Secrets needed: `GROQ_API_KEY` (console.groq.com, free, no credit card)."
        )

        with gr.Row():

            # ── Left: controls ───────────────────────────────────────────
            with gr.Column(scale=1, min_width=300):

                gr.Markdown("## 1. Configure your game")

                game_type_dropdown = gr.Dropdown(
                    choices=GAME_TYPE_NAMES, value="Platformer", label="Game genre",
                )
                type_description = gr.Markdown(
                    value="_" + GAME_TYPES["Platformer"]["description"] + "_",
                )
                theme_box = gr.Textbox(
                    label="Theme / setting",
                    placeholder="e.g. Ancient Egyptian pyramid with cursed mummies",
                    lines=3,
                    value=THEME_EXAMPLES["Platformer"][0][0],
                )
                gr.Examples(
                    examples=THEME_EXAMPLES["Platformer"],
                    inputs=[theme_box],
                    label="Theme examples",
                )

                gr.Markdown("## 2. Generation settings")

                temperature_slider = gr.Slider(
                    minimum=0.3, maximum=1.2, value=0.7, step=0.05,
                    label="Temperature - higher = more creative",
                )
                max_tokens_slider = gr.Slider(
                    minimum=1000, maximum=6000, value=4000, step=500,
                    label="Max tokens - more = longer game",
                )

                generate_btn = gr.Button("Generate Game + Sprites", variant="primary")
                gen_status   = gr.Markdown(value="_No game generated yet._")

            # ── Right: code + game ────────────────────────────────────────
            with gr.Column(scale=2, min_width=500):

                # Collapsible code window
                with gr.Accordion("3. Generated code — click to expand/hide", open=False):
                    code_box = gr.Code(
                        label="HTML source (sprites embedded as base64)",
                        language="html",
                        lines=12,
                        interactive=True,
                    )
                    launch_btn = gr.Button("Launch Game", variant="secondary")

                # Collapsible image prompts window
                with gr.Accordion("3b. Generated image prompts — click to expand/hide", open=False):
                    prompt_display = gr.Markdown(
                        value="_Image prompts will appear here after generation._"
                    )

                gr.Markdown("## 4. Live game window")

                game_frame = gr.HTML(
                    value=_placeholder_html("Generate a game to see it here."),
                )

        # ── Wiring ────────────────────────────────────────────────────────

        game_type_dropdown.change(
            fn=update_type_description,
            inputs=[game_type_dropdown],
            outputs=[type_description],
        )
        game_type_dropdown.change(
            fn=get_first_theme,
            inputs=[game_type_dropdown],
            outputs=[theme_box],
        )

        generate_btn.click(
            fn=generate_game_code,
            inputs=[game_type_dropdown, theme_box, temperature_slider, max_tokens_slider],
            outputs=[code_box, prompt_display, gen_status, game_frame],
        )

        launch_btn.click(fn=launch_game, inputs=[code_box], outputs=[game_frame])

        gr.Markdown(
            "---\n"
            "**Pipeline:** Theme → [Groq Llama] game code + [Groq 70B as Z-Image-Engineer] "
            "cinematic prompts → [FLUX.1-dev/fal-ai] sprites → embedded in game. "
            "Game window fixed to 800×450px — matches canvas exactly. "
            "Edit HTML and click **Launch Game** to hot-reload."
        )

    return demo


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    app = build_ui()
    app.launch()