Update app.py

app.py

@@ -1,21 +1,19 @@
 """
-Text-to-Game Generator - Groq + FLUX Edition
-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
+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-schnell via nscale] --> sprite images
+         --> injected as base64 into game HTML
 
 Secrets needed:
-  GROQ_API_KEY  - from console.groq.com (free)
-  HF_TOKEN      - from huggingface.co/settings/tokens (free)
+  GROQ_API_KEY  - console.groq.com (free, no credit card)
+  HF_TOKEN      - huggingface.co/settings/tokens (for FLUX via nscale)
 """
 
 import os
 import re
 import io
-import json
 import base64
 import traceback
 
@@ -28,28 +26,50 @@ from PIL import Image
 # Clients
 # ---------------------------------------------------------------------------
 
-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", "")
 
+CODE_MODEL  = "llama-3.1-8b-instant"   # Groq — game code
+PROMPT_MODEL = "llama-3.3-70b-versatile" # Groq — better model for creative prompts
+IMAGE_MODEL = "black-forest-labs/FLUX.1-schnell"  # nscale
+
 
-def get_groq():
+def get_groq_client():
     if not GROQ_API_KEY:
         raise ValueError(
-            "GROQ_API_KEY not set.\n"
-            "Get a free key at console.groq.com then add it as a Space secret."
+            "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)
+    return OpenAI(
+        base_url="https://api.groq.com/openai/v1",
+        api_key=GROQ_API_KEY,
+    )
 
 
-def get_flux():
+def get_image_client():
     if not HF_TOKEN:
         return None
     return InferenceClient(provider="nscale", api_key=HF_TOKEN)
 
 
+# ---------------------------------------------------------------------------
+# 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
 # ---------------------------------------------------------------------------
@@ -57,47 +77,102 @@ def get_flux():
 GAME_TYPES = {
     "Platformer": {
         "description": "Jump over enemies and obstacles to reach the goal.",
-        "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"
+        "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"
+            "- Use new Image() with src='sprite_player.png' for the player character.\n"
+            "- Use new Image() with src='sprite_background.png' for the background.\n"
+            "- Use new Image() with src='sprite_enemy.png' for enemies.\n"
+            "- Draw images with ctx.drawImage(img, x, y, w, h).\n"
+            "- Draw platforms and UI with canvas shapes.\n"
+            "- NO external libraries, NO CDN links.\n"
+            "- Theme colors and labels to match: {theme}.\n"
+            "Output ONLY the raw HTML. No explanation, no markdown fences."
         ),
     },
     "Top-Down Shooter": {
         "description": "Shoot waves of enemies before they reach you.",
-        "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"
+        "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"
             "- Game-over screen with score and restart button.\n"
+            "- Use requestAnimationFrame for the game loop.\n"
+            "- Use new Image() with src='sprite_player.png' for the player.\n"
+            "- Use new Image() with src='sprite_background.png' for the background.\n"
+            "- Use new Image() with src='sprite_enemy.png' for enemies.\n"
+            "- Draw images with ctx.drawImage(img, x, y, w, h).\n"
+            "- NO external libraries, NO CDN links.\n"
+            "- Theme colors and labels to match: {theme}.\n"
+            "Output ONLY the raw HTML. No explanation, no markdown fences."
         ),
     },
     "Puzzle / Maze": {
         "description": "Navigate a maze or solve a tile puzzle to escape.",
-        "genre_rules": (
+        "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"
             "- Player navigates with arrow keys or WASD.\n"
-            "- At least 15x10 tile grid with walls, collectible keys, and a locked exit.\n"
-            "- Show a move counter or timer on canvas.\n"
+            "- At least 15x10 tile grid, collectible keys, and a locked exit.\n"
+            "- Show a move counter or timer.\n"
             "- Win screen when player collects all keys and exits.\n"
+            "- Use new Image() with src='sprite_player.png' for the player.\n"
+            "- Use new Image() with src='sprite_background.png' for the background.\n"
+            "- Draw tiles and UI with canvas shapes.\n"
+            "- NO external libraries, NO CDN links.\n"
+            "- Theme colors and labels to match: {theme}.\n"
+            "Output ONLY the raw HTML. No explanation, no markdown fences."
         ),
     },
     "Arcade / Dodge": {
         "description": "Dodge falling obstacles and survive as long as possible.",
-        "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"
+        "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"
+            "- Use new Image() with src='sprite_player.png' for the player.\n"
+            "- Use new Image() with src='sprite_background.png' for the background.\n"
+            "- Draw obstacles and UI with canvas shapes.\n"
+            "- NO external libraries, NO CDN links.\n"
+            "- Theme colors and labels to match: {theme}.\n"
+            "Output ONLY the raw HTML. No explanation, no markdown fences."
         ),
     },
     "Surprise Me!": {
-        "description": "Let the AI invent the genre.",
-        "genre_rules": (
-            "- Pick any fun arcade genre: breakout, snake, flappy-style, space invaders, etc.\n"
+        "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"
             "- Clear win/lose conditions and a score display.\n"
-            "- Game-over or win screen with a restart button.\n"
+            "- Game-over / win screen with a restart button.\n"
+            "- Use requestAnimationFrame for the game loop.\n"
+            "- Use new Image() with src='sprite_player.png' for the player.\n"
+            "- Use new Image() with src='sprite_background.png' for the background.\n"
+            "- Draw other elements with canvas shapes.\n"
+            "- NO external libraries, NO CDN links.\n"
+            "- Theme colors and labels vividly to match: {theme}.\n"
+            "Output ONLY the raw HTML. No explanation, no markdown fences."
         ),
     },
 }
@@ -105,7 +180,7 @@ GAME_TYPES = {
 GAME_TYPE_NAMES = list(GAME_TYPES.keys())
 
 THEME_EXAMPLES = {
-    "Platformer":       [["Jungle temple with cursed mummies"], ["Neon cyberpunk rooftops"], ["Underwater pirate shipwreck"]],
+    "Platformer":       [["Jungle temple with ancient traps"], ["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"]],
@@ -113,58 +188,50 @@ THEME_EXAMPLES = {
 }
 
 # ---------------------------------------------------------------------------
-# Step 1 — Extract sprite prompts from game idea via Groq
+# Step 1: Generate image prompts via Z-Image-Engineer (Groq)
 # ---------------------------------------------------------------------------
 
-def extract_sprite_prompts(game_idea: str, game_type: str) -> list:
+def generate_image_prompts(theme: str, game_type: str) -> dict:
     """
-    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}, ...]
+    Use Groq with Z-Image-Engineer system prompt to generate
+    cinematic image prompts for player and background sprites.
+    Returns dict: { 'sprite_player.png': prompt, 'sprite_background.png': prompt, 'sprite_enemy.png': 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"},
-        ]
-
+    client = get_groq_client()
+
+    seeds = {
+        "sprite_player.png":     f"pixel-art game player character for a {theme} themed {game_type} game, front-facing sprite, vibrant colors, clear silhouette, 64x64 pixel style",
+        "sprite_background.png": f"2D game background scene for a {theme} themed {game_type} game, wide landscape, atmospheric, game art style, 800x450",
+        "sprite_enemy.png":      f"pixel-art enemy character for a {theme} themed {game_type} game, menacing, clear silhouette, 64x64 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  # fallback to raw seed
+    return prompts
 
 # ---------------------------------------------------------------------------
-# Step 2 — Generate sprites with FLUX
+# Step 2: Generate images via FLUX.1-schnell
 # ---------------------------------------------------------------------------
 
-def _pil_to_data_uri(pil_image: Image.Image, size: int = 64) -> str:
-    img = pil_image.resize((size, size), Image.LANCZOS)
+def _pil_to_data_uri(pil_image: Image.Image, size: tuple = None) -> str:
+    if size:
+        img = pil_image.resize(size, Image.LANCZOS)
+    else:
+        img = pil_image
     buf = io.BytesIO()
     img.save(buf, format="PNG")
     return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
@@ -173,7 +240,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[:6]
+    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"/>'
@@ -183,155 +250,118 @@ def _colored_placeholder(name: str) -> str:
     return "data:image/svg+xml;base64," + base64.b64encode(svg.encode()).decode()
 
 
-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 generate_sprites(image_prompts: dict) -> dict:
+    """Generate actual images from enhanced prompts via FLUX.1-schnell."""
+    image_client = get_image_client()
+    sprite_map = {}
+
+    for sprite_name, prompt in image_prompts.items():
+        if not image_client:
+            sprite_map[sprite_name] = _colored_placeholder(sprite_name.replace(".png", ""))
+            continue
+        try:
+            # Background gets full size, sprites get 64x64
+            is_bg = "background" in sprite_name
+            pil_img = image_client.text_to_image(prompt, model=IMAGE_MODEL)
+            size    = None if is_bg else (64, 64)
+            sprite_map[sprite_name] = _pil_to_data_uri(pil_img, size=size)
+            print(f"[FLUX] Generated {sprite_name}")
+        except Exception as exc:
+            print(f"[FLUX] Failed {sprite_name}: {exc}")
+            sprite_map[sprite_name] = _colored_placeholder(sprite_name.replace(".png", ""))
+
+    return sprite_map
 
 # ---------------------------------------------------------------------------
-# Step 3 — Generate game code with sprites already embedded
+# Step 3: Inject sprites into HTML
 # ---------------------------------------------------------------------------
 
-def build_sprite_css(sprite_map: dict) -> str:
-    """Build JS that preloads all sprites as Image objects with base64 src."""
-    lines = ["const sprites = {};"]
+def _inject_sprites(html_code: str, sprite_map: dict) -> str:
     for fname, data_uri in sprite_map.items():
-        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_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
-
+        html_code = html_code.replace(f'"{fname}"', f'"{data_uri}"')
+        html_code = html_code.replace(f"'{fname}'", f"'{data_uri}'")
+    return html_code
 
 # ---------------------------------------------------------------------------
-# Main generation pipeline
+# Step 4: Generate game code via Groq Llama
 # ---------------------------------------------------------------------------
 
-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."""
-
+CODE_SYSTEM = (
+    "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: window.addEventListener('keydown', e => keys.add(e.key)). "
+    "2. Gravity: every frame do player.velY += 0.5 AFTER moving the player. "
+    "3. Platform collision: set grounded=false BEFORE the platform loop. "
+    "   Inside the loop, if player lands on a platform set grounded=true and velY=0. "
+    "4. Jump: if (grounded && (keys.has('ArrowUp') || keys.has('w') || keys.has('W') || keys.has(' '))) { velY=-12; grounded=false; } "
+    "5. Always include a solid ground platform spanning the full canvas width at y=420. "
+    "6. Load sprites: 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'; "
+    "7. Draw background: ctx.drawImage(bgImg, 0, 0, canvas.width, canvas.height); "
+    "8. Draw player/enemies using ctx.drawImage(img, x, y, w, h). "
+    "Output ONLY the raw HTML - no markdown fences, no explanation, nothing else."
+)
+
+
+def generate_game_code(game_type: str, theme: str, temperature: float, max_new_tokens: int):
     if not theme.strip():
-        yield chat_history + [["", "⚠️ Please enter a theme first."]], "", _placeholder_html("Enter a theme.")
-        return
+        return "", "", "Please enter a theme first.", _placeholder_html("Enter a theme and generate a game.")
 
-    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]
+    try:
+        client       = get_groq_client()
+        user_prompt  = GAME_TYPES[game_type]["prompt_template"].format(theme=theme.strip())
 
-    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...")
+        # -- Code generation (Llama) --
+        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 prompt generation (Z-Image-Engineer via Groq) --
+        image_prompts = generate_image_prompts(theme.strip(), game_type)
+
+        # -- Sprite generation (FLUX.1-schnell) --
+        sprite_map = generate_sprites(image_prompts)
+
+        # -- Inject sprites --
+        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
+        status = (
+            f"Done! {n_real} sprite(s) generated by FLUX, "
+            f"{n_fallback} fallback. Click Launch Game to play."
+        )
 
-    # --- Stage 2: Generate sprite images ---
-    history.append(["", "🖼️ Generating sprites with FLUX.1-schnell..."])
-    yield history, "", _placeholder_html("Generating sprites...")
+        # Show the enhanced prompts that were used
+        prompt_summary = "\n\n".join(
+            f"**{k}:**\n{v}" for k, v in image_prompts.items()
+        )
 
-    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...")
+        return final_code, prompt_summary, status, _build_preview(final_code)
 
-    # --- 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 = f"❌ Error generating code: {exc}"
-        history.append(["", err])
-        yield history, "", _placeholder_html(err)
+        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)
 
 
 # ---------------------------------------------------------------------------
@@ -378,7 +408,7 @@ def launch_game(code: str) -> str:
 # UI helpers
 # ---------------------------------------------------------------------------
 
-def update_description(game_type: str) -> str:
+def update_type_description(game_type: str) -> str:
     return "_" + GAME_TYPES[game_type]["description"] + "_"
 
 
@@ -394,10 +424,11 @@ def build_ui():
     with gr.Blocks(title="Game Generator") as demo:
 
         gr.Markdown(
-            "# 🎮 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."
+            "# 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-schnell** renders the sprites.\n\n"
+            "> Secrets needed: `GROQ_API_KEY` (console.groq.com, free) "
+            "and `HF_TOKEN` (huggingface.co/settings/tokens, for FLUX images)."
         )
 
         with gr.Row():
@@ -405,16 +436,16 @@ def build_ui():
             # ── Left: controls ───────────────────────────────────────────
             with gr.Column(scale=1, min_width=300):
 
-                gr.Markdown("## 1. Configure")
+                gr.Markdown("## 1. Configure your game")
 
-                game_type_dd = gr.Dropdown(
+                game_type_dropdown = gr.Dropdown(
                     choices=GAME_TYPE_NAMES, value="Platformer", label="Game genre",
                 )
-                type_desc = gr.Markdown(
-                    value="_" + GAME_TYPES["Platformer"]["description"] + "_"
+                type_description = gr.Markdown(
+                    value="_" + GAME_TYPES["Platformer"]["description"] + "_",
                 )
                 theme_box = gr.Textbox(
-                    label="Game theme / idea",
+                    label="Theme / setting",
                     placeholder="e.g. Ancient Egyptian pyramid with cursed mummies",
                     lines=3,
                     value=THEME_EXAMPLES["Platformer"][0][0],
@@ -425,40 +456,38 @@ def build_ui():
                     label="Theme examples",
                 )
 
-                gr.Markdown("## 2. Settings")
+                gr.Markdown("## 2. Generation settings")
 
-                temp_slider = gr.Slider(
+                temperature_slider = gr.Slider(
                     minimum=0.3, maximum=1.2, value=0.7, step=0.05,
-                    label="Temperature",
+                    label="Temperature - higher = more creative",
                 )
-                tokens_slider = gr.Slider(
+                max_tokens_slider = gr.Slider(
                     minimum=1000, maximum=6000, value=4000, step=500,
-                    label="Max tokens",
+                    label="Max tokens - more = longer game",
                 )
 
-                generate_btn = gr.Button("🚀 Generate Game", variant="primary")
+                generate_btn = gr.Button("Generate Game + Sprites", variant="primary")
+                gen_status   = gr.Markdown(value="_No game generated yet._")
 
-            # ── Right: chat + code + game ─────────────────────────────────
-            with gr.Column(scale=2, min_width=500):
-
-                gr.Markdown("## 3. Generation log")
-
-                chatbox = gr.Chatbot(
-                    label="Sprite & code pipeline",
-                    height=180,
-                    show_label=False,
+                gr.Markdown("## 3. Generated image prompts")
+                prompt_display = gr.Markdown(
+                    value="_Image prompts will appear here after generation._"
                 )
 
-                gr.Markdown("## 4. Generated code _(editable)_")
+            # ── Right: code + game ────────────────────────────────────────
+            with gr.Column(scale=2, min_width=500):
+
+                gr.Markdown("## 4. Generated code (editable)")
 
                 code_box = gr.Code(
-                    label="HTML source",
+                    label="HTML source (sprites embedded as base64)",
                     language="html",
-                    lines=10,
+                    lines=12,
                     interactive=True,
                 )
 
-                launch_btn = gr.Button("▶ Launch Game", variant="secondary")
+                launch_btn = gr.Button("Launch Game", variant="secondary")
 
                 gr.Markdown("## 5. Live game window")
 
@@ -468,22 +497,30 @@ def build_ui():
 
         # ── 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])
+        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=run_generation,
-            inputs=[game_type_dd, theme_box, temp_slider, tokens_slider, chatbox],
-            outputs=[chatbox, code_box, game_frame],
+            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 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."
+            "**Pipeline:** Theme → [Groq Llama] game code + [Groq 70B as Z-Image-Engineer] "
+            "cinematic prompts → [FLUX.1-schnell/nscale] sprites → embedded in game. "
+            "Edit the HTML and click **Launch Game** to hot-reload."
         )
 
     return demo