api.py

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from backend.environment import GameConfig
from backend.engine import Engine
from backend.tools import TOOL_SCHEMAS

app = FastAPI(title="Grid Royale API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

engine: Engine | None = None


class StartGameRequest(BaseModel):
    grid_size: int = 20
    max_turns: int = 50
    max_agents: int = 8
    num_chests: int = 15
    agent_names: list[str] | None = None
    system_prompt: str | None = None
    system_prompts: dict[str, str] | None = None


def _serialize(aid: str) -> dict:
    if not engine:
        return {}
    a = engine.env.agents[aid]
    return {
        "id": a.id,
        "name": a.name,
        "x": a.pos[0],
        "y": a.pos[1],
        "hp": a.hp,
        "alive": a.alive,
        "score": a.score,
        "kills": a.kills,
        "shielded": a.shielded,
        "abilities": [
            {"name": ab.name, "last_used_turn": ab.last_used_turn, "uses_remaining": ab.uses_remaining}
            for ab in a.abilities
        ],
        "messages": a.messages[-24:],
    }


def _game_state() -> dict:
    if not engine:
        return {}
    alive = engine.env.alive_agents()
    tiles = {}
    for (x, y), tile in engine.env.grid.items():
        tiles[f"{x},{y}"] = {
            "terrain": tile.terrain,
            "loot": [
                {
                    "type": item.type,
                    "ability_name": item.ability_name,
                    "points": item.points,
                    "rarity": item.rarity,
                    "source": item.source,
                }
                for item in tile.loot
            ] if tile.loot else [],
        }
    return {
        "turn": engine.env.turn,
        "grid_size": engine.env.config.grid_size,
        "max_turns": engine.env.config.max_turns,
        "agents": [_serialize(aid) for aid in engine.env.agents],
        "tiles": tiles,
        "game_over": len(alive) <= 1 or engine.env.turn >= engine.env.config.max_turns,
        "winner": alive[0].name if len(alive) == 1 else None,
        "turn_log": engine.last_turn_log,
        "tool_schemas": TOOL_SCHEMAS,
    }


BOT_NAMES = [
    "Vex", "Nyx", "Kael", "Zara",
    "Cypher", "Helix", "Apex", "Blaze",
]


def _default_names(count: int) -> list[str]:
    return BOT_NAMES[:count]


@app.post("/api/game/start")
async def start_game(req: StartGameRequest):
    global engine
    config = GameConfig(
        grid_size=req.grid_size,
        max_turns=req.max_turns,
        max_agents=req.max_agents,
        num_chests=req.num_chests,
    )
    engine = Engine(config=config)
    engine.generate_grid()
    engine.scatter_chests()

    names = req.agent_names or _default_names(req.max_agents)
    for i, name in enumerate(names):
        if i >= req.max_agents:
            break
        aid = f"agent_{i}"
        prompt = None
        if req.system_prompts:
            prompt = req.system_prompts.get(aid) or req.system_prompts.get(name)
        if not prompt:
            prompt = req.system_prompt
        engine.spawn_agent(agent_id=aid, name=name, system_prompt=prompt)

    return {"message": f"Game started with {len(names)} agents", "turn": 0}


@app.post("/api/game/step")
async def game_step():
    if not engine:
        raise HTTPException(400, "No game running")
    await engine.step()
    return _game_state()


@app.get("/api/game/state")
async def game_state():
    if not engine:
        raise HTTPException(400, "No game running")
    return _game_state()


@app.post("/api/game/run")
async def run_full_game(req: StartGameRequest):
    global engine
    config = GameConfig(
        grid_size=req.grid_size,
        max_turns=req.max_turns,
        max_agents=req.max_agents,
        num_chests=req.num_chests,
    )
    engine = Engine(config=config)
    engine.generate_grid()
    engine.scatter_chests()

    names = req.agent_names or _default_names(req.max_agents)
    for i, name in enumerate(names):
        if i >= req.max_agents:
            break
        aid = f"agent_{i}"
        prompt = None
        if req.system_prompts:
            prompt = req.system_prompts.get(aid) or req.system_prompts.get(name)
        if not prompt:
            prompt = req.system_prompt
        engine.spawn_agent(agent_id=aid, name=name, system_prompt=prompt)

    await engine.run_game()
    return _game_state()


@app.get("/api/game/default_prompt")
def get_default_prompt(grid_size: int = 20):
    from backend.agent.agent import Agent
    return {"default_prompt": Agent.get_default_system_prompt(grid_size=grid_size)}


@app.post("/api/model/warm")
async def warm_model():
    from backend.agent.llm_client import get_llm_client
    try:
        client = get_llm_client()
        response = await client.chat.completions.create(
            model="google/gemma-4-26B-A4B-it",
            messages=[{"role": "user", "content": "ping"}],
            max_tokens=5
        )
        if response and response.choices:
            return {
                "status": "success",
                "message": "Model warmed up successfully!",
                "response": response.choices[0].message.content
            }
        return {"status": "warning", "message": "No choices returned from model."}
    except Exception as e:
        return {"status": "error", "detail": str(e)}