agent.py

"""Ollama-backed game generation agent using structured JSON output."""
from __future__ import annotations

import copy
import json
import re

import ollama
from pydantic import ValidationError

from schema import (
    CanvasGame, CollisionAction, InputAction, InputKey, MovementPattern, Shape,
)

# A canonical example game baked into the prompt — the single strongest signal
# for the model to follow the schema structure correctly.
_EXAMPLE = json.dumps({
    "title": "Neon Dodge",
    "width": 600,
    "height": 500,
    "background_color": "#0a0a0a",
    "sprites": [
        {
            "id": "player", "shape": "triangle", "color": "#00ffcc",
            "x": 300, "y": 440, "width": 32, "height": 32,
            "movement_pattern": "keyboard_horizontal",
            "speed": 5.0, "is_player": True, "count": 1,
        },
        {
            "id": "rock", "shape": "circle", "color": "#ff4444",
            "x": 100, "y": 0, "width": 26, "height": 26,
            "movement_pattern": "rain_down",
            "speed": 4.0, "is_player": False, "count": 8,
        },
    ],
    "collision_rules": [
        {"between": ["player", "rock"], "action": "game_over", "value": 1},
    ],
    "input_bindings": [
        {"key": "arrow_left",  "action": "move_left",  "target": "player"},
        {"key": "arrow_right", "action": "move_right", "target": "player"},
    ],
    "score_display": {
        "position": "top_right", "font_size": 22,
        "color": "#00ffcc", "label": "SCORE",
    },
    "lives": 3,
}, indent=2)

_SYSTEM_PROMPT = f"""/no_think
You are a creative arcade game designer. Output a single JSON object.
The JSON must start with {{ and end with }} — no wrapper key, no markdown fences, no extra text.

REQUIRED top-level keys (use exactly these names):
  "title", "width", "height", "background_color",
  "sprites"          ← NON-EMPTY list; must contain a player and at least one enemy or item
  "collision_rules"  ← NON-EMPTY list
  "input_bindings"   ← NON-EMPTY list whose target matches the player sprite id
  "score_display", "lives"

MOVEMENT PATTERNS:
  keyboard_horizontal — player moves left/right with arrow keys  ← use for most players
  keyboard_all        — player moves in all 4 directions
  rain_down           — falls from top, respawns (enemies/coins)
  bounce              — bounces off all walls
  chase_player        — homes toward the player
  zigzag              — side-to-side drift downward
  circular            — orbits a fixed point
  static              — doesn't move

VALID VALUES (use only these — no others):
  shape:             circle | rectangle | triangle | star
  movement_pattern:  keyboard_horizontal | keyboard_all | rain_down | bounce |
                     chase_player | zigzag | circular | static
  collision action:  game_over | score_plus | score_minus | destroy_target | bounce_back
  input key:         arrow_left | arrow_right | arrow_up | arrow_down | space
  input action:      move_left | move_right | move_up | move_down   (no "shoot")
  score position:    top_left | top_right | top_center

RULES:
- Exactly ONE sprite with is_player=true; use keyboard_horizontal or keyboard_all
- input_bindings MUST reference the player sprite id exactly
- Use count=6-10 for enemies/collectibles so the screen feels alive
- Colors: vivid neon hex — #0ff, #f0f, #ff0, #0f0, #f44, #f80, #88f, #0f8, etc.
- game_over action when player touches an enemy
- score_plus when player touches a collectible; destroy_target to remove an enemy on touch
- Title: 2-4 evocative words, e.g. "Void Escape", "Star Blitz", "Neon Rain"

EXAMPLE (follow this structure exactly):
{_EXAMPLE}
"""


_VALID_SHAPES      = {s.value for s in Shape}
_VALID_MOVEMENTS   = {m.value for m in MovementPattern}
_VALID_INPUT_KEYS  = {k.value for k in InputKey}
_VALID_INPUT_ACTS  = {a.value for a in InputAction}
_VALID_COL_ACTS    = {a.value for a in CollisionAction}


_FIELD_ALIASES: dict[str, str] = {
    "game_title":       "title",
    "name":             "title",
    "background":       "background_color",
    "bg_color":         "background_color",
    "backgroundColor":  "background_color",
    "sprite_list":      "sprites",
    "sprite_objects":   "sprites",
    "controls":         "input_bindings",
    "inputs":           "input_bindings",
    "key_bindings":     "input_bindings",
    "collisions":       "collision_rules",
    "collision_list":   "collision_rules",
    "hud":              "score_display",
    "display":          "score_display",
    "max_lives":        "lives",
    "num_lives":        "lives",
}


def _coerce(data: dict) -> dict:
    """Normalise field names, clamp values, and fix invalid enums before Pydantic validation."""
    data = copy.deepcopy(data)

    # Rename aliased top-level keys
    for alias, canonical in _FIELD_ALIASES.items():
        if alias in data and canonical not in data:
            data[canonical] = data.pop(alias)

    for s in data.get("sprites", []):
        if s.get("shape") not in _VALID_SHAPES:
            s["shape"] = "circle"
        if s.get("movement_pattern") not in _VALID_MOVEMENTS:
            s["movement_pattern"] = "static"
        s["speed"]  = max(0.5, min(15.0, float(s.get("speed",  3.0))))
        s["count"]  = max(1,   min(20,   int(  s.get("count",  1  ))))
        s["width"]  = max(5,   min(150,  float(s.get("width",  30 ))))
        s["height"] = max(5,   min(150,  float(s.get("height", 30 ))))

    data["input_bindings"] = [
        b for b in data.get("input_bindings", [])
        if b.get("key")    in _VALID_INPUT_KEYS
        and b.get("action") in _VALID_INPUT_ACTS
    ]

    clean_rules = []
    for r in data.get("collision_rules", []):
        if r.get("action") not in _VALID_COL_ACTS:
            r["action"] = "game_over"
        between = r.get("between", [])
        if isinstance(between, list) and len(between) >= 2:
            r["between"] = between[:2]
            clean_rules.append(r)
    data["collision_rules"] = clean_rules

    data["width"]  = max(300, min(800, int(data.get("width",  600))))
    data["height"] = max(300, min(650, int(data.get("height", 500))))

    return data


class GameAgent:
    def __init__(
        self,
        model: str = "qwen3.5:9b",
        ollama_host: str = "http://127.0.0.1:11434",
    ) -> None:
        self.model = model
        self.client = ollama.Client(host=ollama_host)
        self.current_game_json: dict | None = None

    def _generate(self, messages: list[dict], max_retries: int = 2) -> CanvasGame:
        last_error: Exception | None = None

        for attempt in range(max_retries + 1):
            response = self.client.chat(
                model=self.model,
                messages=messages,
                format="json",
                options={"temperature": 0.7 + attempt * 0.05, "num_predict": 3000},
            )
            content: str = response.message.content

            # Qwen3 may emit <think>…</think> before the JSON
            if "</think>" in content:
                content = content.split("</think>")[-1].strip()

            try:
                data = json.loads(content)
            except json.JSONDecodeError:
                # Extract first {...} block from surrounding prose
                match = re.search(r"\{.*\}", content, re.DOTALL)
                if not match:
                    last_error = ValueError(f"No JSON found in response: {content[:200]}")
                    continue
                try:
                    data = json.loads(match.group())
                except json.JSONDecodeError as e:
                    last_error = e
                    continue

            # Unwrap common model wrapping patterns like {"game": {...}}
            for _key in ("game", "config", "output", "result", "data", "arcade_game"):
                if isinstance(data, dict) and _key in data and isinstance(data[_key], dict):
                    data = data[_key]
                    break

            data = _coerce(data)

            try:
                game = CanvasGame.model_validate(data)
            except ValidationError as e:
                last_error = e
                continue

            # Quality gate: must have at least one player sprite
            if not any(s.is_player for s in game.sprites):
                last_error = ValueError(
                    "Model did not include a player sprite (is_player=true). "
                    "Retrying with stronger hint."
                )
                # Inject a stronger hint for the next attempt
                messages = _inject_player_hint(messages)
                continue

            self.current_game_json = game.model_dump()
            return game

        raise ValueError(
            f"Failed to generate a valid game after {max_retries + 1} attempts. "
            f"Last error: {last_error}"
        )

    def create_game(self, prompt: str) -> CanvasGame:
        messages = [
            {"role": "system", "content": _SYSTEM_PROMPT},
            {
                "role": "user",
                "content": (
                    f"Create a game: {prompt}\n\n"
                    "Output a single JSON object with these exact top-level keys: "
                    "title, width, height, background_color, sprites, collision_rules, "
                    "input_bindings, score_display, lives."
                ),
            },
        ]
        return self._generate(messages)

    def modify_game(self, modification: str) -> CanvasGame:
        if self.current_game_json is None:
            raise ValueError("No active game to modify. Create a game first.")
        current = json.dumps(self.current_game_json, indent=2)
        messages = [
            {"role": "system", "content": _SYSTEM_PROMPT},
            {
                "role": "user",
                "content": (
                    f"Current game JSON:\n{current}\n\n"
                    f"Modify it: {modification}\n\n"
                    "Output the complete updated JSON object."
                ),
            },
        ]
        return self._generate(messages)

    def reset(self) -> None:
        self.current_game_json = None


def _inject_player_hint(messages: list[dict]) -> list[dict]:
    """Append a reminder that a player sprite is required."""
    messages = list(messages)
    messages.append({
        "role": "assistant",
        "content": "{}",  # placeholder so next user turn makes sense
    })
    messages.append({
        "role": "user",
        "content": (
            "IMPORTANT: your sprites array must include exactly one sprite "
            "with \"is_player\": true and movement_pattern \"keyboard_horizontal\" "
            "or \"keyboard_all\". Please regenerate the complete JSON."
        ),
    })
    return messages