HuggingWizards-upload/app.py

"""HuggingWizards — co-op pixel wizard arena with a Nemotron-4B Game Master.

Single process:
  * FastAPI serves the canvas client (`/`) and a WebSocket (`/ws`).
  * A ~20 Hz asyncio task runs the authoritative simulation on CPU and
    broadcasts snapshots to every connected client (players + spectators).
  * At each round boundary the Game Master (Nemotron, burst GPU) decides rewards
    and the next round; every decision is logged as an agent trace.
  * A Gradio dashboard is mounted at `/dashboard` (live traces + leaderboard).
"""
from __future__ import annotations

import asyncio
import contextlib
import json
import os
import time
import uuid

import gradio as gr
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

from game import gamemaster
from game.engine import Engine, MAX_PLAYERS, TICK_HZ

HERE = os.path.dirname(__file__)
engine = Engine()

# Active websocket connections: ws -> {"pid", "role", "name", "char"}
clients: dict[WebSocket, dict] = {}
# Ordered waiting line for when the 8-slot arena is full.
queue: list[WebSocket] = []

# ---- persistent leaderboard ----------------------------------------------
LEADERBOARD_FILE = os.path.join(HERE, "traces", "leaderboard.json")
_leaderboard: list[dict] = []


def _load_leaderboard():
    global _leaderboard
    try:
        with open(LEADERBOARD_FILE, encoding="utf-8") as f:
            _leaderboard = json.load(f)
    except Exception:
        _leaderboard = []


def _save_leaderboard():
    with contextlib.suppress(Exception):
        with open(LEADERBOARD_FILE, "w", encoding="utf-8") as f:
            json.dump(_leaderboard[:100], f, indent=2)


def add_score(name: str, gold: int, level: int, wave: int) -> list[dict]:
    _leaderboard.append({
        "name": str(name)[:16] or "Wizard", "gold": int(gold),
        "level": int(level), "wave": int(wave),
        "ts": time.strftime("%Y-%m-%d %H:%M"),
    })
    _leaderboard.sort(key=lambda e: (e["gold"], e["wave"], e["level"]), reverse=True)
    del _leaderboard[100:]
    _save_leaderboard()
    return _leaderboard[:20]


_load_leaderboard()


async def broadcast(payload: dict):
    if not clients:
        return
    msg = json.dumps(payload)
    dead = []
    for ws in list(clients.keys()):
        try:
            await ws.send_text(msg)
        except Exception:
            dead.append(ws)
    for ws in dead:
        _drop(ws)


def _drop(ws: WebSocket):
    info = clients.pop(ws, None)
    if info and info.get("pid"):
        engine.remove_player(info["pid"])
        if not engine.players and engine.status != "lobby":
            engine.status = "lobby"
            engine.status_msg = "All wizards left. Waiting for players..."


async def _send(ws: WebSocket, payload: dict):
    with contextlib.suppress(Exception):
        await ws.send_text(json.dumps(payload))


async def handle_skill_request(ws: WebSocket, pid: str, prompt: str):
    """The round's top wizard asks the Game Master for a custom power-up."""
    if engine.status != "intermission" or engine.skill_request_used \
            or engine.skill_request_pid != pid or not prompt.strip():
        await _send(ws, {"type": "skill_result", "ok": False,
                         "reason": "Not available right now."})
        return
    engine.skill_request_used = True  # one wish per intermission
    ctx = {"round": engine.round, "player_level": engine.players[pid].level
           if pid in engine.players else 1}
    spec = await asyncio.to_thread(gamemaster.generate_skill, prompt, ctx)
    granted = engine.add_runtime_skill(pid, spec) if spec else None
    await _send(ws, {"type": "skill_result", "ok": bool(granted), "skill": granted,
                     "reason": None if granted else "The arcane energies fizzled — try again."})


def _ws_for_pid(pid: str):
    for ws, info in clients.items():
        if info.get("pid") == pid:
            return ws
    return None


async def _broadcast_queue():
    """Tell each queued connection its current place in line."""
    for i, ws in enumerate(list(queue)):
        await _send(ws, {"type": "queue", "pos": i + 1, "size": len(queue)})


async def promote_queue():
    """Fill open arena slots from the front of the queue."""
    while engine.can_join() and queue:
        ws = queue.pop(0)
        info = clients.get(ws)
        if info is None:      # disconnected while waiting
            continue
        pid = uuid.uuid4().hex[:8]
        engine.add_player(pid, info.get("name", "Wizard"), info.get("char", "warrior"))
        clients[ws] = {**info, "pid": pid, "role": "player"}
        await _send(ws, {"type": "welcome", "pid": pid, "role": "player", "reason": "from_queue"})


async def manage_slots():
    """Eliminate out-of-lives wizards, then promote the queue into free slots."""
    for pid, p in list(engine.players.items()):
        if p.lives <= 0 and not p.alive:
            ws = _ws_for_pid(pid)
            engine.remove_player(pid)
            if ws is not None:
                info = clients.get(ws, {})
                clients[ws] = {**info, "pid": None, "role": "spectator"}
                await _send(ws, {"type": "eliminated",
                                 "score": {"gold": p.gold, "level": p.level, "wave": engine.round}})
    if not engine.players and engine.status != "lobby":
        engine.status = "lobby"
        engine.status_msg = "Waiting for wizards to join..."
    await promote_queue()
    await _broadcast_queue()


async def game_loop():
    last = time.time()
    prev_status = engine.status
    gm_task: asyncio.Task | None = None
    gm_applied = False
    interval = 1.0 / TICK_HZ
    slot_accum = 0.0
    while True:
        now = time.time()
        dt = min(0.1, now - last)
        last = now

        engine.tick(dt)
        # eliminate / promote queue / broadcast positions a few times a second
        slot_accum += dt
        if slot_accum >= 0.4:
            slot_accum = 0.0
            await manage_slots()

        # Round just ended -> ask the Game Master (off the loop thread).
        if engine.status == "intermission" and prev_status == "active":
            summary = engine.round_summary()
            prev_cfg = dict(engine.cfg)
            gm_task = asyncio.create_task(
                asyncio.to_thread(gamemaster.decide, summary, prev_cfg)
            )
            gm_applied = False

        # Apply the decision as soon as it's ready (enables shopping this break).
        if gm_task is not None and gm_task.done() and not gm_applied:
            with contextlib.suppress(Exception):
                engine.apply_gm_decision(gm_task.result())
            gm_applied = True

        # Start the next round once the break is over AND the decision is in.
        if engine.status == "intermission" and gm_applied and time.time() >= engine.intermission_until:
            if engine.players:
                engine.start_round()
            else:
                engine.status = "lobby"
                engine.status_msg = "Waiting for wizards to join..."
            gm_task = None
            gm_applied = False

        prev_status = engine.status
        await broadcast(engine.snapshot())
        await asyncio.sleep(max(0, interval - (time.time() - now)))


# --------------------------------------------------------------------------
# FastAPI app
# --------------------------------------------------------------------------
app = FastAPI(title="HuggingWizards")

_loop_task: asyncio.Task | None = None


def ensure_loop():
    """Start the game loop once. Robust to Gradio overriding the app lifespan."""
    global _loop_task
    if _loop_task is None or _loop_task.done():
        _loop_task = asyncio.create_task(game_loop())


@app.get("/")
async def index():
    return FileResponse(os.path.join(HERE, "static", "index.html"))


@app.get("/traces")
async def traces():
    """Recent Game Master agent traces, trimmed for the in-game panel."""
    out = []
    for t in gamemaster.RECENT_TRACES[:15]:
        out.append({
            "trace_id": t.get("trace_id"), "round": t.get("round"),
            "ts": t.get("ts"), "model": t.get("model"),
            "source": t.get("source"), "latency_sec": t.get("latency_sec"),
            "kind": t.get("kind", "round_decision"), "error": t.get("error"),
            "mercy": t.get("mercy"),
            "raw_output": str(t.get("raw_output", ""))[:600],
            "decision": t.get("decision"),
        })
    return {"traces": out}


@app.websocket("/ws")
async def ws_endpoint(ws: WebSocket):
    await ws.accept()
    ensure_loop()
    clients[ws] = {"pid": None, "role": "spectator", "name": "Wizard", "char": "warrior"}
    try:
        while True:
            data = json.loads(await ws.receive_text())
            mtype = data.get("type")

            if mtype == "join":
                name = str(data.get("name", "Wizard"))[:16] or "Wizard"
                char = str(data.get("char", "warrior"))[:16] or "soldier"
                clients[ws]["name"] = name
                clients[ws]["char"] = char
                if clients[ws]["pid"]:
                    continue
                if engine.can_join():
                    pid = uuid.uuid4().hex[:8]
                    engine.add_player(pid, name, char)
                    clients[ws] = {**clients[ws], "pid": pid, "role": "player"}
                    await _send(ws, {"type": "welcome", "pid": pid, "role": "player"})
                else:
                    # game full -> join the queue
                    if ws not in queue:
                        queue.append(ws)
                    pos = queue.index(ws) + 1
                    await _send(ws, {"type": "welcome", "pid": None, "role": "spectator",
                                     "reason": "queued", "pos": pos, "size": len(queue)})

            elif mtype == "spectate":
                clients[ws]["name"] = str(data.get("name", "Wizard"))[:16] or "Wizard"
                clients[ws]["char"] = str(data.get("char", "warrior"))[:16] or "soldier"

            elif mtype == "leave_queue":
                if ws in queue:
                    queue.remove(ws)

            elif mtype == "get_leaderboard":
                await _send(ws, {"type": "leaderboard", "top": _leaderboard[:20]})

            elif mtype == "submit_score":
                top = add_score(str(data.get("name", "Wizard")), int(data.get("gold", 0)),
                                int(data.get("level", 1)), int(data.get("wave", 0)))
                await _send(ws, {"type": "leaderboard", "top": top, "submitted": True})

            elif mtype == "input" and clients[ws]["pid"]:
                engine.set_input(clients[ws]["pid"], data)

            elif mtype == "choose_card" and clients[ws]["pid"]:
                engine.choose_card(clients[ws]["pid"], data.get("key", ""))

            elif mtype == "request_skill" and clients[ws]["pid"]:
                await handle_skill_request(ws, clients[ws]["pid"], str(data.get("prompt", "")))

            elif mtype == "start" and clients[ws]["pid"]:
                if engine.status == "lobby":
                    engine.start_round()
    except WebSocketDisconnect:
        pass
    except Exception:
        pass
    finally:
        if ws in queue:
            queue.remove(ws)
        _drop(ws)


app.mount("/static", StaticFiles(directory=os.path.join(HERE, "static")), name="static")


# --------------------------------------------------------------------------
# Gradio dashboard (mounted at /dashboard)
# --------------------------------------------------------------------------
def _dashboard_state():
    players = sorted(engine.players.values(), key=lambda p: p.gold, reverse=True)
    if players:
        lb = "| Wizard | Level | Gold | Last-round dmg | Kills |\n|---|---|---|---|---|\n"
        lb += "\n".join(
            f"| {p.name} | {p.level} | {p.gold} | {round(p.dmg_dealt)} | {p.kills} |"
            for p in players
        )
    else:
        lb = "_No wizards in the arena yet._"
    header = (f"### Round {engine.round} — `{engine.status}`\n"
              f"{engine.status_msg}\n\n"
              f"Players: {len(engine.players)}/{MAX_PLAYERS} · "
              f"GM: _{engine.gm_message or '—'}_")
    traces = gamemaster.RECENT_TRACES[:10]
    return header, lb, traces


with gr.Blocks(title="HuggingWizards — Game Master Dashboard") as demo:
    gr.Markdown("# 🧙 HuggingWizards — Game Master Dashboard\n"
                "Play at the [arena](/). Below: live state, leaderboard, and the "
                "per-round **agent traces** produced by Nemotron-4B.")
    header_md = gr.Markdown()
    with gr.Row():
        leaderboard_md = gr.Markdown()
    gr.Markdown("## Agent traces (most recent first)")
    traces_json = gr.JSON()
    timer = gr.Timer(2.0)
    timer.tick(_dashboard_state, outputs=[header_md, leaderboard_md, traces_json])
    demo.load(_dashboard_state, outputs=[header_md, leaderboard_md, traces_json])

app = gr.mount_gradio_app(app, demo, path="/dashboard")


if __name__ == "__main__":
    import uvicorn

    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)