Update app.py

app.py

@@ -1,14 +1,21 @@
 """
 Text-to-Game Generator - Groq + FLUX Edition
-- Code  : llama-3.1-8b-instant via Groq API (free, no credit card)
-- Images: black-forest-labs/FLUX.1-schnell via nscale (HF Inference API, free with HF_TOKEN)
-
-Setup: add GROQ_API_KEY and HF_TOKEN as Space secrets.
+Flow:
+  1. User types game idea in chat box
+  2. Groq extracts sprite names + image prompts from the idea
+  3. FLUX.1-schnell generates each sprite image
+  4. Groq generates the full game code using the actual sprite data URIs
+  5. Game runs in browser iframe with real sprites embedded
+
+Secrets needed:
+  GROQ_API_KEY  - from console.groq.com (free)
+  HF_TOKEN      - from huggingface.co/settings/tokens (free)
 """
 
 import os
 import re
 import io
+import json
 import base64
 import traceback
 
@@ -18,136 +25,79 @@ from huggingface_hub import InferenceClient
 from PIL import Image
 
 # ---------------------------------------------------------------------------
-# Google AI Studio client (OpenAI-compatible endpoint)
+# Clients
 # ---------------------------------------------------------------------------
 
-CODE_MODEL    = "llama-3.1-8b-instant"
-GROQ_API_KEY  = os.environ.get("GROQ_API_KEY", "")
+CODE_MODEL  = "llama-3.1-8b-instant"
+IMAGE_MODEL = "black-forest-labs/FLUX.1-schnell"
+
+GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
+HF_TOKEN     = os.environ.get("HF_TOKEN", "")
 
 
-def _check_key():
+def get_groq():
     if not GROQ_API_KEY:
         raise ValueError(
-            "GROQ_API_KEY is not set.\n"
-            "1. Go to console.groq.com and sign in\n"
-            "2. Click API Keys > Create API Key\n"
-            "3. Copy the key (starts with gsk_)\n"
-            "4. Add it as a Space secret: Settings > Variables and secrets > Name: GROQ_API_KEY"
+            "GROQ_API_KEY not set.\n"
+            "Get a free key at console.groq.com then add it as a Space secret."
         )
+    return OpenAI(base_url="https://api.groq.com/openai/v1", api_key=GROQ_API_KEY)
 
 
-def get_client():
-    _check_key()
-    return OpenAI(
-        base_url="https://api.groq.com/openai/v1",
-        api_key=GROQ_API_KEY,
-    )
-
-
-# ---------------------------------------------------------------------------
-# Image model — FLUX.1-schnell via nscale (free with HF_TOKEN)
-# ---------------------------------------------------------------------------
-
-IMAGE_MODEL = "black-forest-labs/FLUX.1-schnell"
-HF_TOKEN    = os.environ.get("HF_TOKEN", "")
-
-
-def get_image_client():
+def get_flux():
     if not HF_TOKEN:
         return None
     return InferenceClient(provider="nscale", api_key=HF_TOKEN)
 
 
 # ---------------------------------------------------------------------------
-# Game type configs — uses sprite_NAME.png when HF_TOKEN set, else canvas
+# Game type configs
 # ---------------------------------------------------------------------------
 
 GAME_TYPES = {
     "Platformer": {
         "description": "Jump over enemies and obstacles to reach the goal.",
-        "prompt_template": (
-            "Create a complete, self-contained HTML5 platformer game with the theme: {theme}.\n"
-            "Requirements:\n"
-            "- Single HTML file with all CSS and JavaScript inline.\n"
-            "- Use an HTML5 canvas (id='gameCanvas') sized 800x450.\n"
-            "- Player moves left/right with arrow keys or WASD and jumps.\n"
-            "- At least 5 platforms, 2 moving enemies, and a goal to reach.\n"
-            "- Show score/lives on the canvas.\n"
-            "- Use requestAnimationFrame for the game loop.\n"
-            "- Draw ALL graphics using canvas 2D shapes and colors (fillRect, arc, fillStyle, etc).\n"
-            "- Make colors, shapes and labels vividly match the theme: {theme}.\n"
-            "- NO image files, NO external libraries, NO CDN links.\n"
-            "Output ONLY the raw HTML. No explanation, no markdown fences."
+        "genre_rules": (
+            "- Player moves left/right with arrow keys or WASD and jumps with ArrowUp/W/Space.\n"
+            "- At least 5 platforms including a full-width ground at y=420.\n"
+            "- 2 moving enemies that patrol platforms.\n"
+            "- A goal object the player must reach to win.\n"
+            "- Show score and lives on canvas.\n"
         ),
     },
     "Top-Down Shooter": {
         "description": "Shoot waves of enemies before they reach you.",
-        "prompt_template": (
-            "Create a complete, self-contained HTML5 top-down shooter game with the theme: {theme}.\n"
-            "Requirements:\n"
-            "- Single HTML file with all CSS and JavaScript inline.\n"
-            "- Use an HTML5 canvas (id='gameCanvas') sized 800x450.\n"
-            "- Player moves with WASD/arrows; shoots with Space or click.\n"
-            "- Enemies spawn from edges in escalating waves.\n"
-            "- Display health, score, and wave on canvas.\n"
+        "genre_rules": (
+            "- Player moves with WASD/arrows; shoots with Space or click toward mouse.\n"
+            "- Enemies spawn from canvas edges in escalating waves.\n"
+            "- Display health, score, and wave number on canvas.\n"
             "- Game-over screen with score and restart button.\n"
-            "- Use requestAnimationFrame for the game loop.\n"
-            "- Draw ALL graphics using canvas 2D shapes and colors only.\n"
-            "- Make colors, shapes and labels vividly match the theme: {theme}.\n"
-            "- NO image files, NO external libraries, NO CDN links.\n"
-            "Output ONLY the raw HTML. No explanation, no markdown fences."
         ),
     },
     "Puzzle / Maze": {
         "description": "Navigate a maze or solve a tile puzzle to escape.",
-        "prompt_template": (
-            "Create a complete, self-contained HTML5 maze/tile-puzzle game with the theme: {theme}.\n"
-            "Requirements:\n"
-            "- Single HTML file with all CSS and JavaScript inline.\n"
-            "- Use an HTML5 canvas (id='gameCanvas') sized 800x450.\n"
+        "genre_rules": (
             "- Player navigates with arrow keys or WASD.\n"
-            "- At least 15x10 tile grid, collectible keys, and a locked exit.\n"
-            "- Show a move counter or timer.\n"
+            "- At least 15x10 tile grid with walls, collectible keys, and a locked exit.\n"
+            "- Show a move counter or timer on canvas.\n"
             "- Win screen when player collects all keys and exits.\n"
-            "- Draw ALL graphics using canvas 2D shapes and colors only.\n"
-            "- Make colors, shapes and labels vividly match the theme: {theme}.\n"
-            "- NO image files, NO external libraries, NO CDN links.\n"
-            "Output ONLY the raw HTML. No explanation, no markdown fences."
         ),
     },
     "Arcade / Dodge": {
         "description": "Dodge falling obstacles and survive as long as possible.",
-        "prompt_template": (
-            "Create a complete, self-contained HTML5 arcade dodge game with the theme: {theme}.\n"
-            "Requirements:\n"
-            "- Single HTML file with all CSS and JavaScript inline.\n"
-            "- Use an HTML5 canvas (id='gameCanvas') sized 800x450.\n"
-            "- Player moves with arrow keys or WASD; obstacles increase in speed over time.\n"
-            "- Collision ends the game; show time survived as score.\n"
-            "- High score stored in a JS variable; restart button on game-over screen.\n"
-            "- Use requestAnimationFrame for the game loop.\n"
-            "- Draw ALL graphics using canvas 2D shapes and colors only.\n"
-            "- Make colors, shapes and labels vividly match the theme: {theme}.\n"
-            "- NO image files, NO external libraries, NO CDN links.\n"
-            "Output ONLY the raw HTML. No explanation, no markdown fences."
+        "genre_rules": (
+            "- Player moves left/right with arrow keys or WASD.\n"
+            "- Obstacles fall from the top, increasing in speed over time.\n"
+            "- Show time survived as score.\n"
+            "- Game-over screen with high score and restart button.\n"
         ),
     },
     "Surprise Me!": {
-        "description": "Let the AI invent the genre - could be anything!",
-        "prompt_template": (
-            "Create a complete, self-contained HTML5 browser game with the theme: {theme}.\n"
-            "Pick any fun arcade genre: breakout, snake, flappy-style, space invaders, etc.\n"
-            "Requirements:\n"
-            "- Single HTML file with all CSS and JavaScript inline.\n"
-            "- Use an HTML5 canvas (id='gameCanvas') sized 800x450.\n"
-            "- Keyboard or mouse controlled.\n"
+        "description": "Let the AI invent the genre.",
+        "genre_rules": (
+            "- Pick any fun arcade genre: breakout, snake, flappy-style, space invaders, etc.\n"
             "- Clear win/lose conditions and a score display.\n"
-            "- Game-over / win screen with a restart button.\n"
-            "- Use requestAnimationFrame for the game loop.\n"
-            "- Draw ALL graphics using canvas 2D shapes and colors only.\n"
-            "- Make colors, shapes and labels vividly match the theme: {theme}.\n"
-            "- NO image files, NO external libraries, NO CDN links.\n"
-            "Output ONLY the raw HTML. No explanation, no markdown fences."
+            "- Game-over or win screen with a restart button.\n"
         ),
     },
 }
@@ -155,7 +105,7 @@ GAME_TYPES = {
 GAME_TYPE_NAMES = list(GAME_TYPES.keys())
 
 THEME_EXAMPLES = {
-    "Platformer":       [["Jungle temple with ancient traps"], ["Neon cyberpunk rooftops"], ["Underwater pirate shipwreck"]],
+    "Platformer":       [["Jungle temple with cursed mummies"], ["Neon cyberpunk rooftops"], ["Underwater pirate shipwreck"]],
     "Top-Down Shooter": [["Alien desert invasion"], ["Viking village under siege"], ["Steampunk robot uprising"]],
     "Puzzle / Maze":    [["Haunted library with secret doors"], ["Ice cave with frozen keys"], ["Egyptian pyramid tiles"]],
     "Arcade / Dodge":   [["Asteroid field in a tiny rocket"], ["Chef dodging ingredients"], ["Time traveller avoiding paradox storms"]],
@@ -163,7 +113,54 @@ THEME_EXAMPLES = {
 }
 
 # ---------------------------------------------------------------------------
-# Sprite helpers
+# Step 1 — Extract sprite prompts from game idea via Groq
+# ---------------------------------------------------------------------------
+
+def extract_sprite_prompts(game_idea: str, game_type: str) -> list:
+    """
+    Ask Groq to decide what 3 sprites the game needs and write
+    a detailed FLUX image prompt for each one.
+    Returns a list of dicts: [{name, filename, flux_prompt}, ...]
+    """
+    system = (
+        "You are a game art director. Given a game idea, decide what 3 sprite images "
+        "are needed and write a detailed image generation prompt for each. "
+        "Output ONLY valid JSON — a list of 3 objects, each with keys: "
+        "'name' (short label e.g. 'player'), "
+        "'filename' (e.g. 'sprite_player.png'), "
+        "'flux_prompt' (detailed prompt for FLUX image model, pixel-art style, 64x64, dark background). "
+        "No explanation, no markdown, just the JSON array."
+    )
+    user = f"Game type: {game_type}\nGame idea: {game_idea}"
+
+    try:
+        client = get_groq()
+        resp = client.chat.completions.create(
+            model=CODE_MODEL,
+            messages=[
+                {"role": "system", "content": system},
+                {"role": "user",   "content": user},
+            ],
+            max_tokens=500,
+            temperature=0.5,
+        )
+        raw = resp.choices[0].message.content.strip()
+        raw = re.sub(r"^```[a-zA-Z]*\n?", "", raw).strip()
+        raw = re.sub(r"\n?```$", "", raw).strip()
+        sprites = json.loads(raw)
+        return sprites[:3]
+    except Exception as exc:
+        print("[Sprite extractor] Error:", exc)
+        # Fallback: generic sprites
+        return [
+            {"name": "player",     "filename": "sprite_player.png",     "flux_prompt": f"Pixel-art player sprite for a {game_idea} game, 64x64, dark background"},
+            {"name": "enemy",      "filename": "sprite_enemy.png",       "flux_prompt": f"Pixel-art enemy sprite for a {game_idea} game, 64x64, dark background"},
+            {"name": "background", "filename": "sprite_background.png",  "flux_prompt": f"Pixel-art background tile for a {game_idea} game, 64x64, dark background"},
+        ]
+
+
+# ---------------------------------------------------------------------------
+# Step 2 — Generate sprites with FLUX
 # ---------------------------------------------------------------------------
 
 def _pil_to_data_uri(pil_image: Image.Image, size: int = 64) -> str:
@@ -176,7 +173,7 @@ def _pil_to_data_uri(pil_image: Image.Image, size: int = 64) -> str:
 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]
+    label   = name[:6]
     svg = (
         '<svg xmlns="http://www.w3.org/2000/svg" width="64" height="64">'
         '<rect width="64" height="64" fill="' + colour + '" rx="8"/>'
@@ -186,130 +183,155 @@ def _colored_placeholder(name: str) -> str:
     return "data:image/svg+xml;base64," + base64.b64encode(svg.encode()).decode()
 
 
-def _sprite_prompt(sprite_name: str, theme: str) -> str:
-    label = sprite_name.replace("sprite_", "").replace("_", " ")
-    return (
-        "Pixel-art game sprite: " + label + ". Theme: " + theme + ". "
-        "Vibrant colours, clear silhouette, plain dark background, 64x64 pixel style."
-    )
-
+def generate_sprite_images(sprite_list: list) -> dict:
+    """
+    Given list of {name, filename, flux_prompt}, generate each image.
+    Returns dict: {filename: data_uri}
+    """
+    client = get_flux()
+    mapping = {}
+    for sprite in sprite_list:
+        fname  = sprite["filename"]
+        prompt = sprite["flux_prompt"]
+        name   = sprite["name"]
+        if client:
+            try:
+                pil_img = client.text_to_image(prompt, model=IMAGE_MODEL)
+                mapping[fname] = _pil_to_data_uri(pil_img, size=64)
+                print(f"[FLUX] Generated {fname}")
+            except Exception as exc:
+                print(f"[FLUX] Failed {fname}: {exc}")
+                mapping[fname] = _colored_placeholder(name)
+        else:
+            mapping[fname] = _colored_placeholder(name)
+    return mapping
 
-def _find_sprite_filenames(html_code: str) -> list:
-    pattern = r"""(?:src\s*=\s*['"]|\.src\s*=\s*['"])\s*(sprite_[a-zA-Z0-9_]+\.png)\s*['"]"""
-    return list(dict.fromkeys(re.findall(pattern, html_code)))[:3]
 
+# ---------------------------------------------------------------------------
+# Step 3 — Generate game code with sprites already embedded
+# ---------------------------------------------------------------------------
 
-def _inject_sprites(html_code: str, sprite_map: dict) -> str:
+def build_sprite_css(sprite_map: dict) -> str:
+    """Build JS that preloads all sprites as Image objects with base64 src."""
+    lines = ["const sprites = {};"]
     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
+        var_name = fname.replace(".png", "").replace("-", "_")
+        lines.append(f'sprites["{fname}"] = new Image();')
+        lines.append(f'sprites["{fname}"].src = "{data_uri}";')
+    return "\n".join(lines)
 
 
-def generate_sprites(sprite_names: list, theme: str) -> dict:
-    if not sprite_names:
-        return {}
-    client = get_image_client()
-    if not client:
-        return {n: _colored_placeholder(n.replace(".png","")) for n in sprite_names}
-    mapping = {}
-    for fname in sprite_names:
-        name_no_ext = fname.replace(".png", "")
-        prompt = _sprite_prompt(name_no_ext, theme)
-        try:
-            pil_img = client.text_to_image(prompt, model=IMAGE_MODEL)
-            mapping[fname] = _pil_to_data_uri(pil_img, size=64)
-        except Exception as exc:
-            print("[FLUX] Failed '" + fname + "': " + str(exc))
-            mapping[fname] = _colored_placeholder(name_no_ext)
-    return mapping
+def generate_game_code(game_type: str, theme: str, temperature: float,
+                       max_new_tokens: int, sprite_map: dict) -> str:
+    """Generate HTML game code, injecting sprite data URIs directly into the prompt."""
+
+    genre_rules = GAME_TYPES[game_type]["genre_rules"]
+
+    # Build sprite instruction with actual filenames
+    sprite_filenames = list(sprite_map.keys())
+    sprite_instruction = (
+        "Available sprites (already loaded as Image objects in a 'sprites' object):\n"
+        + "\n".join(f'  sprites["{f}"] — use as: ctx.drawImage(sprites["{f}"], x, y, w, h)' for f in sprite_filenames)
+        + "\n"
+        "Use ALL of these sprites to draw game objects. Do NOT use fillRect for objects that have a sprite.\n"
+    )
+
+    # Build sprite preload JS to inject into the HTML
+    sprite_preload_js = build_sprite_css(sprite_map)
+
+    system_msg = (
+        "You are an expert HTML5 game developer. "
+        "Write a complete, working, single-file HTML5 game using canvas and vanilla JavaScript. "
+        "CRITICAL RULES — follow every one exactly:\n"
+        "1. Add keydown/keyup listeners on WINDOW: window.addEventListener('keydown', e => keys.add(e.key)).\n"
+        "2. Use a Set called 'keys' for input tracking.\n"
+        "3. Gravity: every frame add 0.5 to velY AFTER position update.\n"
+        "4. Platform collision: set grounded=false BEFORE the platform loop. "
+        "Inside the loop set grounded=true and velY=0 when player lands on top.\n"
+        "5. Jump: if (grounded && (keys.has('ArrowUp')||keys.has('w')||keys.has('W')||keys.has(' '))) { velY=-12; grounded=false; }\n"
+        "6. Always include a full-width ground platform at y=420.\n"
+        "7. The sprites are PRE-LOADED — include this exact JS block near the top of your script:\n"
+        f"```\n{sprite_preload_js}\n```\n"
+        "8. Use ctx.drawImage(sprites['filename.png'], x, y, w, h) to draw sprites.\n"
+        "9. Output ONLY the raw HTML — no markdown fences, no explanation."
+    )
+
+    user_msg = (
+        f"Create a complete HTML5 {game_type} game with the theme: {theme}.\n"
+        "Requirements:\n"
+        "- Single HTML file, all CSS and JS inline.\n"
+        "- Canvas id='gameCanvas' sized 800x450.\n"
+        f"{genre_rules}"
+        "- Use requestAnimationFrame for the game loop.\n"
+        f"{sprite_instruction}"
+        "Output ONLY the raw HTML."
+    )
+
+    client = get_groq()
+    resp = client.chat.completions.create(
+        model=CODE_MODEL,
+        messages=[
+            {"role": "system", "content": system_msg},
+            {"role": "user",   "content": user_msg},
+        ],
+        max_tokens=int(max_new_tokens),
+        temperature=float(temperature),
+    )
+    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)
+    return code
 
 
 # ---------------------------------------------------------------------------
-# Code generation
+# Main generation pipeline
 # ---------------------------------------------------------------------------
 
-def generate_game_code(game_type: str, theme: str, temperature: float, max_new_tokens: int):
+def run_generation(game_type: str, theme: str, temperature: float,
+                   max_new_tokens: int, chat_history: list):
+    """Full pipeline: extract sprites → generate images → generate game code."""
+
     if not theme.strip():
-        return "", "Please enter a theme first.", _placeholder_html("Enter a theme and generate a game.")
-
-    template    = GAME_TYPES[game_type]["prompt_template"]
-    user_prompt = template.format(theme=theme.strip())
-    use_sprites = bool(HF_TOKEN)
-
-    if use_sprites:
-        system_msg = (
-            "You are an expert HTML5 game developer. "
-            "Write a complete, working, single-file HTML5 game using canvas and vanilla JavaScript. "
-            "CRITICAL RULES - follow every one exactly: "
-            "1. Add keydown/keyup listeners on WINDOW not canvas: window.addEventListener('keydown', e => keys.add(e.key)). "
-            "2. Gravity: every frame do velY += 0.5 AFTER moving, NOT inside collision. "
-            "3. Platform collision: loop ALL platforms, find the one the player is standing on. "
-               "   Set grounded=true and velY=0 ONLY when found. Set grounded=false BEFORE the loop, not inside else. "
-            "4. Jump: if (grounded && (keys.has('ArrowUp') || keys.has('w') || keys.has('W'))) { velY=-12; grounded=false; } "
-            "5. Always add a solid ground platform spanning the full canvas width at y=420. "
-            "6. Use new Image() with src='sprite_NAME.png' for every visual asset, max 3 sprites. "
-            "7. Draw images with ctx.drawImage(img, x, y, w, h). "
-            "Output ONLY the raw HTML - no markdown fences, no explanation."
-        )
-    else:
-        system_msg = (
-            "You are an expert HTML5 game developer. "
-            "Write a complete, working, single-file HTML5 game using only canvas 2D drawing. "
-            "CRITICAL RULES - follow every one exactly: "
-            "1. Add keydown/keyup listeners on WINDOW not canvas: window.addEventListener('keydown', e => keys.add(e.key)). "
-            "2. Gravity: every frame do velY += 0.5 AFTER moving, NOT inside collision. "
-            "3. Platform collision: loop ALL platforms, find the one the player is standing on. "
-               "   Set grounded=true and velY=0 ONLY when found. Set grounded=false BEFORE the loop, not inside else. "
-            "4. Jump: if (grounded && (keys.has('ArrowUp') || keys.has('w') || keys.has('W'))) { velY=-12; grounded=false; } "
-            "5. Always add a solid ground platform spanning the full canvas width at y=420. "
-            "6. Never use image files. Draw everything with canvas shapes and colors. "
-            "Output ONLY the raw HTML - no markdown fences, no explanation."
-        )
+        yield chat_history + [["", "⚠️ Please enter a theme first."]], "", _placeholder_html("Enter a theme.")
+        return
 
-    try:
-        client = get_client()
-        completion = client.chat.completions.create(
-            model=CODE_MODEL,
-            messages=[
-                {"role": "system", "content": system_msg},
-                {"role": "user",   "content": user_prompt},
-            ],
-            max_tokens=int(max_new_tokens),
-            temperature=float(temperature),
-        )
-        code = completion.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)
-
-        if use_sprites:
-            sprite_names = _find_sprite_filenames(code)
-            sprite_map   = generate_sprites(sprite_names, theme.strip())
-            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
-            status = (
-                "Done! " + str(len(sprite_names)) + " sprite(s) - "
-                + str(n_real) + " by FLUX.1-schnell, "
-                + str(n_fallback) + " fallback. Click Launch Game to play."
-            )
-        else:
-            status = "Done! Game generated (canvas graphics). Click Launch Game to play."
+    history = list(chat_history)
+
+    # --- Stage 1: Extract sprite prompts ---
+    history.append([theme, "🎨 Deciding what sprites your game needs..."])
+    yield history, "", _placeholder_html("Analysing your game idea...")
+
+    sprites_meta = extract_sprite_prompts(theme, game_type)
+    sprite_names = [s["filename"] for s in sprites_meta]
+
+    sprite_summary = "**Sprites planned:**\n" + "\n".join(
+        f"- `{s['filename']}` — {s['name']}: _{s['flux_prompt'][:80]}..._"
+        for s in sprites_meta
+    )
+    history[-1][1] = sprite_summary
+    yield history, "", _placeholder_html("Generating sprite images with FLUX...")
+
+    # --- Stage 2: Generate sprite images ---
+    history.append(["", "🖼️ Generating sprites with FLUX.1-schnell..."])
+    yield history, "", _placeholder_html("Generating sprites...")
 
-        return code, status, _build_preview(code)
+    sprite_map = generate_sprite_images(sprites_meta)
+    n_real = sum(1 for v in sprite_map.values() if "image/png" in v)
+    history[-1][1] = f"✅ {n_real}/{len(sprites_meta)} sprites generated by FLUX. Now writing game code..."
+    yield history, "", _placeholder_html("Writing game code with Groq...")
 
+    # --- Stage 3: Generate game code ---
+    try:
+        code = generate_game_code(game_type, theme, temperature, max_new_tokens, sprite_map)
+        history.append(["", f"✅ Game ready! ({len(code.splitlines())} lines of code)"])
+        yield history, code, _build_preview(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)
+        err = f"❌ Error generating code: {exc}"
+        history.append(["", err])
+        yield history, "", _placeholder_html(err)
 
 
 # ---------------------------------------------------------------------------
@@ -356,7 +378,7 @@ def launch_game(code: str) -> str:
 # UI helpers
 # ---------------------------------------------------------------------------
 
-def update_type_description(game_type: str) -> str:
+def update_description(game_type: str) -> str:
     return "_" + GAME_TYPES[game_type]["description"] + "_"
 
 
@@ -369,33 +391,32 @@ def get_first_theme(game_type: str) -> str:
 # ---------------------------------------------------------------------------
 
 def build_ui():
-    with gr.Blocks(title="Game Generator - Free") as demo:
+    with gr.Blocks(title="Game Generator") as demo:
 
         gr.Markdown(
-            "# Game Generator (Free)\n"
-            "**Gemini 2.5 Flash** via Google AI Studio generates complete playable games.\n"
-            "Canvas-drawn graphics — no image API, no quota, no cost.\n\n"
-            "> **Setup:** Go to [ai.google.dev](https://ai.google.dev), sign in with any Google account, "
-            "click **Get API key** (free, no credit card, 2 minutes). "
-            "Then add it as a Space secret named `GOOGLE_API_KEY`."
+            "# 🎮 Game Generator\n"
+            "Type your game idea below. The AI will plan the sprites, "
+            "generate them with **FLUX.1-schnell**, then write the full game with **Llama-3.1-8B**.\n\n"
+            "> Needs `GROQ_API_KEY` and `HF_TOKEN` as Space secrets."
         )
 
         with gr.Row():
 
+            # ── Left: controls ───────────────────────────────────────────
             with gr.Column(scale=1, min_width=300):
 
-                gr.Markdown("## 1. Configure your game")
+                gr.Markdown("## 1. Configure")
 
-                game_type_dropdown = gr.Dropdown(
+                game_type_dd = gr.Dropdown(
                     choices=GAME_TYPE_NAMES, value="Platformer", label="Game genre",
                 )
-                type_description = gr.Markdown(
-                    value="_" + GAME_TYPES["Platformer"]["description"] + "_",
+                type_desc = gr.Markdown(
+                    value="_" + GAME_TYPES["Platformer"]["description"] + "_"
                 )
                 theme_box = gr.Textbox(
-                    label="Theme / setting",
-                    placeholder="e.g. Neon cyberpunk rooftops",
-                    lines=2,
+                    label="Game theme / idea",
+                    placeholder="e.g. Ancient Egyptian pyramid with cursed mummies",
+                    lines=3,
                     value=THEME_EXAMPLES["Platformer"][0][0],
                 )
                 gr.Examples(
@@ -404,56 +425,65 @@ def build_ui():
                     label="Theme examples",
                 )
 
-                gr.Markdown("## 2. Generation settings")
+                gr.Markdown("## 2. Settings")
 
-                temperature_slider = gr.Slider(
+                temp_slider = gr.Slider(
                     minimum=0.3, maximum=1.2, value=0.7, step=0.05,
-                    label="Temperature - higher = more creative",
+                    label="Temperature",
                 )
-                max_tokens_slider = gr.Slider(
+                tokens_slider = gr.Slider(
                     minimum=1000, maximum=6000, value=4000, step=500,
-                    label="Max tokens - more = longer / more complete game",
+                    label="Max tokens",
                 )
 
-                generate_btn = gr.Button("Generate Game", variant="primary")
-                gen_status   = gr.Markdown(value="_No game generated yet._")
+                generate_btn = gr.Button("🚀 Generate Game", variant="primary")
 
+            # ── Right: chat + code + game ─────────────────────────────────
             with gr.Column(scale=2, min_width=500):
 
-                gr.Markdown("## 3. Generated code (editable)")
+                gr.Markdown("## 3. Generation log")
+
+                chatbox = gr.Chatbot(
+                    label="Sprite & code pipeline",
+                    height=180,
+                    show_label=False,
+                )
+
+                gr.Markdown("## 4. Generated code _(editable)_")
 
                 code_box = gr.Code(
                     label="HTML source",
                     language="html",
-                    lines=12,
+                    lines=10,
                     interactive=True,
                 )
 
-                launch_btn = gr.Button("Launch Game", variant="secondary")
+                launch_btn = gr.Button("▶ Launch Game", variant="secondary")
 
-                gr.Markdown("## 4. Live game window")
+                gr.Markdown("## 5. Live game window")
 
                 game_frame = gr.HTML(
                     value=_placeholder_html("Generate a game to see it here."),
                 )
 
-        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])
+        # ── Wiring ────────────────────────────────────────────────────────
+
+        game_type_dd.change(fn=update_description, inputs=[game_type_dd], outputs=[type_desc])
+        game_type_dd.change(fn=get_first_theme,    inputs=[game_type_dd], outputs=[theme_box])
 
         generate_btn.click(
-            fn=generate_game_code,
-            inputs=[game_type_dropdown, theme_box, temperature_slider, max_tokens_slider],
-            outputs=[code_box, gen_status, game_frame],
+            fn=run_generation,
+            inputs=[game_type_dd, theme_box, temp_slider, tokens_slider, chatbox],
+            outputs=[chatbox, code_box, game_frame],
         )
 
         launch_btn.click(fn=launch_game, inputs=[code_box], outputs=[game_frame])
 
         gr.Markdown(
             "---\n"
-            "**Why Google AI Studio:** 1,500 free requests/day, no credit card, "
-            "Gemini 2.5 Flash is within 5% of GPT-4o quality. "
-            "Rate limits reset daily — not monthly. "
-            "Edit the HTML and click **Launch Game** to hot-reload."
+            "**Pipeline:** Theme → Groq plans sprites → FLUX generates images → "
+            "Groq writes game code with sprites pre-embedded as base64 → "
+            "Game runs in browser iframe. Edit HTML and click **▶ Launch Game** to hot-reload."
         )
 
     return demo