game_agent.py

"""smolagents orchestrator for the spec-driven arcade renderer.

The agent chooses one of a small catalog of render tools and supplies JSON-like
arguments. Tools validate with Pydantic and render deterministic p5.js games.
"""
from __future__ import annotations

import os
from pathlib import Path
from typing import Any

from dotenv import load_dotenv
from pydantic import ValidationError
from smolagents import InferenceClientModel, LiteLLMModel, ToolCallingAgent, tool

from game_engine import GameEngine, GameSpec, render_spec_html

load_dotenv(dotenv_path=Path(__file__).with_name(".env"))

_engine: GameEngine | None = None


def _build_spec(archetype: str, **kwargs: Any) -> GameSpec:
    data = {
        "archetype": archetype,
        "tags": [archetype, "agent-render"],
        **{key: value for key, value in kwargs.items() if value is not None},
    }
    tags = data.get("tags") or []
    if archetype not in tags:
        tags = [archetype, *tags]
    data["tags"] = tags[:5]
    return GameSpec.model_validate(data)


def _render(spec: GameSpec) -> str:
    if _engine is None:
        return "Error: game engine not initialized."
    _engine.current_spec = spec
    _engine.current_html = render_spec_html(spec)
    return (
        f"Rendered {spec.title} with `{spec.archetype}`. "
        f"Goal: {spec.goal}. Controls: {spec.controls}."
    )


@tool
def render_dodge_game(
    title: str,
    theme: str,
    goal: str,
    player_name: str = "runner",
    difficulty: int = 5,
    player_speed: int = 5,
    hazard_count: int = 8,
    hazard_speed: int = 5,
    collectible_count: int = 4,
    goal_score: int = 20,
    twist: str = "falling hazards with glowing pickups",
) -> str:
    """Render a dodge game from validated design fields.

    Args:
        title: Short game title.
        theme: Visual theme and mood.
        goal: Playable objective.
        player_name: Player character noun.
        difficulty: 1-10 difficulty.
        player_speed: 1-10 movement responsiveness.
        hazard_count: Number of hazards, 0-20.
        hazard_speed: 1-10 hazard speed.
        collectible_count: Number of collectibles, 0-20.
        goal_score: Score needed to win, 1-100.
        twist: Strange or delightful flavor detail.
    """
    try:
        spec = _build_spec(
            "dodge",
            title=title,
            theme=theme,
            goal=goal,
            player_name=player_name,
            difficulty=difficulty,
            player_speed=player_speed,
            hazard_count=hazard_count,
            hazard_speed=hazard_speed,
            collectible_count=collectible_count,
            goal_score=goal_score,
            controls="Arrow keys or WASD to move",
            twist=twist,
        )
    except ValidationError as e:
        return f"Validation failed for dodge spec: {e}"
    return _render(spec)


@tool
def render_collector_game(
    title: str,
    theme: str,
    goal: str,
    player_name: str = "collector",
    difficulty: int = 5,
    player_speed: int = 5,
    hazard_count: int = 5,
    hazard_speed: int = 4,
    collectible_count: int = 8,
    goal_score: int = 30,
    twist: str = "collectibles pulse like tiny wishes",
) -> str:
    """Render a collector game from validated design fields.

    Args:
        title: Short game title.
        theme: Visual theme and mood.
        goal: Playable objective.
        player_name: Player character noun.
        difficulty: 1-10 difficulty.
        player_speed: 1-10 movement responsiveness.
        hazard_count: Number of enemies/hazards, 0-20.
        hazard_speed: 1-10 hazard speed.
        collectible_count: Number of collectibles, 0-20.
        goal_score: Score needed to win, 1-100.
        twist: Strange or delightful flavor detail.
    """
    try:
        spec = _build_spec(
            "collector",
            title=title,
            theme=theme,
            goal=goal,
            player_name=player_name,
            difficulty=difficulty,
            player_speed=player_speed,
            hazard_count=hazard_count,
            hazard_speed=hazard_speed,
            collectible_count=collectible_count,
            goal_score=goal_score,
            controls="Arrow keys or WASD to move",
            twist=twist,
        )
    except ValidationError as e:
        return f"Validation failed for collector spec: {e}"
    return _render(spec)


@tool
def render_jumper_game(
    title: str,
    theme: str,
    goal: str,
    player_name: str = "jumper",
    difficulty: int = 5,
    player_speed: int = 5,
    jump_strength: int = 6,
    hazard_count: int = 3,
    hazard_speed: int = 4,
    collectible_count: int = 7,
    goal_score: int = 20,
    twist: str = "reachable platforms arranged like a tiny pilgrimage",
) -> str:
    """Render a platform jumper with validated, reachable platform rules.

    Args:
        title: Short game title.
        theme: Visual theme and mood.
        goal: Playable objective.
        player_name: Player character noun.
        difficulty: 1-10 difficulty.
        player_speed: 1-10 horizontal movement responsiveness.
        jump_strength: 1-10 jump height.
        hazard_count: Number of strange hazards, 0-20.
        hazard_speed: 1-10 hazard speed.
        collectible_count: Number of platforms/collectibles, 0-20.
        goal_score: Score target, 1-100.
        twist: Strange or delightful flavor detail.
    """
    try:
        spec = _build_spec(
            "jumper",
            title=title,
            theme=theme,
            goal=goal,
            player_name=player_name,
            difficulty=difficulty,
            player_speed=player_speed,
            jump_strength=jump_strength,
            hazard_count=hazard_count,
            hazard_speed=hazard_speed,
            collectible_count=collectible_count,
            goal_score=goal_score,
            controls="Left/right to move, Space/W/Up to jump",
            twist=twist,
        )
    except ValidationError as e:
        return f"Validation failed for jumper spec: {e}"
    return _render(spec)


@tool
def render_snake_game(
    title: str,
    theme: str,
    goal: str,
    player_name: str = "snake",
    difficulty: int = 5,
    player_speed: int = 5,
    collectible_count: int = 1,
    goal_score: int = 50,
    twist: str = "direct-control snake trail that only moves on key press",
) -> str:
    """Render a direct-control snake/grid game from validated design fields.

    Args:
        title: Short game title.
        theme: Visual theme and mood.
        goal: Playable objective.
        player_name: Player character noun.
        difficulty: 1-10 difficulty.
        player_speed: 1-10 movement responsiveness.
        collectible_count: Number of active collectibles, 0-20.
        goal_score: Score needed to win, 1-100.
        twist: Strange or delightful flavor detail.
    """
    try:
        spec = _build_spec(
            "snake",
            title=title,
            theme=theme,
            goal=goal,
            player_name=player_name,
            difficulty=difficulty,
            player_speed=player_speed,
            hazard_count=0,
            hazard_speed=3,
            collectible_count=collectible_count,
            goal_score=goal_score,
            controls="Arrow keys or WASD move one grid cell",
            twist=twist,
            direct_control=True,
        )
    except ValidationError as e:
        return f"Validation failed for snake spec: {e}"
    return _render(spec)


@tool
def render_strange_mix_game(
    title: str,
    primary_archetype: str,
    secondary_archetype: str,
    theme: str,
    goal: str,
    player_name: str = "wanderer",
    difficulty: int = 6,
    weirdness: int = 9,
    player_speed: int = 5,
    jump_strength: int = 6,
    hazard_count: int = 7,
    hazard_speed: int = 5,
    collectible_count: int = 7,
    goal_score: int = 30,
    twist: str = "two arcade rules accidentally sharing one dream",
) -> str:
    """Render a strange but playable blend of two render archetypes.

    Use when the user asks for weird, strange, surreal, mixed, cursed, whimsical,
    or surprising gameplay.

    Args:
        title: Short game title.
        primary_archetype: Main render type: dodge, collector, jumper, or snake.
        secondary_archetype: Second render type to blend in.
        theme: Visual theme and mood.
        goal: Playable objective.
        player_name: Player character noun.
        difficulty: 1-10 difficulty.
        weirdness: 1-10 strangeness; use 7-10 for this tool.
        player_speed: 1-10 movement responsiveness.
        jump_strength: 1-10 jump height for jumper mixes.
        hazard_count: Number of hazards, 0-20.
        hazard_speed: 1-10 hazard speed.
        collectible_count: Number of collectibles/platforms, 0-20.
        goal_score: Score needed to win, 1-100.
        twist: Strange or delightful flavor detail.
    """
    try:
        spec = _build_spec(
            primary_archetype,
            title=title,
            secondary_archetype=secondary_archetype,
            weirdness=weirdness,
            tags=[primary_archetype, secondary_archetype, "strange-mix"],
            theme=theme,
            goal=goal,
            player_name=player_name,
            difficulty=difficulty,
            player_speed=player_speed,
            jump_strength=jump_strength,
            hazard_count=hazard_count,
            hazard_speed=hazard_speed,
            collectible_count=collectible_count,
            goal_score=goal_score,
            controls="Arrow keys/WASD; Space jumps in jumper mixes",
            twist=twist,
        )
    except ValidationError as e:
        return f"Validation failed for strange mix spec: {e}"
    return _render(spec)


_AGENT_INSTRUCTIONS = """\
You are the Arcade Director for a Gradio game builder.

You never write JavaScript. You choose exactly one render tool:
- render_dodge_game
- render_collector_game
- render_jumper_game
- render_snake_game
- render_strange_mix_game

Use these tools as the only way to create playable games.
The tool validates and renders the game. If a tool returns a validation error,
fix the arguments and call the tool again.

Choose the render type from the user's idea:
- dodge: avoid falling/projectile hazards
- collector: gather items while avoiding enemies
- jumper: platforms, jumping, reaching a final platform
- snake: grid movement, eating/growing/trails
- strange mix: user asks for weird, strange, cursed, surreal, mixed, dreamlike,
  or surprising; blend two of the four render types

Keep responses short. After a tool succeeds, tell the user the title, render tag,
controls, goal, and strange mix if used.
"""


class GameOrchestrator:
    """Wraps smolagents ToolCallingAgent + deterministic renderer tools."""

    def __init__(
        self,
        model: str = "Qwen/Qwen2.5-Coder-32B-Instruct",
        ollama_host: str = "http://127.0.0.1:11434",
        provider: str | None = None,
    ) -> None:
        global _engine
        provider = (provider or os.getenv("MODEL_PROVIDER", "hf")).lower()
        self.engine = GameEngine(model=model, ollama_host=ollama_host, provider=provider)
        _engine = self.engine

        self.model_id = model
        self.ollama_host = ollama_host
        self.provider = provider
        self._prev_html: str | None = None

        if provider == "ollama":
            llm = LiteLLMModel(
                model_id=f"ollama_chat/{model}",
                api_base=ollama_host,
                api_key="ollama",
                num_ctx=8192,
            )
        else:
            token = (
                os.getenv("HF_TOKEN")
                or os.getenv("HUGGINGFACEHUB_API_TOKEN")
                or os.getenv("HUGGING_TOKEN")
            )
            llm = InferenceClientModel(
                model_id=model,
                token=token,
                timeout=120,
            )

        self.agent = ToolCallingAgent(
            tools=[
                render_dodge_game,
                render_collector_game,
                render_jumper_game,
                render_snake_game,
                render_strange_mix_game,
            ],
            model=llm,
            instructions=_AGENT_INSTRUCTIONS,
            max_steps=4,
            verbosity_level=1,
        )

    def chat(self, message: str) -> tuple[str, str | None]:
        """Process a user message. Returns (reply_text, game_html_or_None)."""
        try:
            result = self.agent.run(message, reset=False)
        except Exception as e:
            return f"Something went wrong while directing the arcade: {e}", None

        html_changed = self.engine.current_html != self._prev_html
        new_html = self.engine.current_html if html_changed else None
        if html_changed:
            self._prev_html = self.engine.current_html

        return str(result), new_html

    def new_game(self) -> None:
        """Reset state for a fresh game session."""
        self.engine.reset()
        self._prev_html = None
        self.agent.memory.reset()

    def has_game(self) -> bool:
        return self.engine.current_html is not None