commander-ai-lab/routes/ws_game.py

"""WebSocket route for real-time game state streaming.

Endpoints:
  WS   /ws/game/{game_id}         — Stream state deltas to Unity
  POST /api/ws/game/{game_id}/push — Push a state update (called by Forge IPC)
  POST /api/ws/game/{game_id}/decision — Push AI decision event
  POST /api/ws/game/{game_id}/end  — Signal game over
  GET  /api/ws/stats               — Connection statistics

Unity connects to the WebSocket endpoint and receives:
  1. An initial full snapshot on first state push
  2. Subsequent delta-only messages as state changes
  3. Decision events when the AI acts
  4. A game_over event when the game ends
"""
from __future__ import annotations

import copy
import logging
import time
import uuid
from typing import Optional

from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from pydantic import BaseModel, Field

from ws.connection_manager import game_manager
from ws.game_state import get_tracker, remove_tracker

log = logging.getLogger("routes.ws_game")

router = APIRouter(tags=["websocket"])


# ────────────────────────────────────────────────────────────
# WebSocket endpoint  —  Unity connects here
# ────────────────────────────────────────────────────────────

@router.websocket("/ws/game/{game_id}")
async def game_ws(websocket: WebSocket, game_id: str):
    """WebSocket endpoint for Unity game-state streaming.

    Query params:
        client_id  —  optional unique identifier for this client

    The connection stays open for the duration of the game.
    The server pushes messages; the client can send:
        {"type": "ping"}            —  keep-alive
        {"type": "request_snapshot"} —  request a full re-sync
    """
    client_id = (
        websocket.query_params.get("client_id")
        or f"unity-{uuid.uuid4().hex[:8]}"
    )

    client = await game_manager.connect(websocket, game_id, client_id)

    # If the tracker already has a baseline, send the current state
    tracker = get_tracker(game_id)
    if tracker.has_baseline:
        snapshot_msg = tracker.update(tracker._last_snapshot)
        await websocket.send_json(snapshot_msg)

    try:
        while True:
            data = await websocket.receive_json()
            msg_type = data.get("type", "")

            if msg_type == "ping":
                await websocket.send_json({"type": "pong"})

            elif msg_type == "request_snapshot":
                if tracker.has_baseline:
                    snap = copy.deepcopy(tracker._last_snapshot)
                    tracker._seq += 1
                    await websocket.send_json({
                        "type": "snapshot",
                        "seq": tracker._seq,
                        "game_id": game_id,
                        "timestamp": time.time(),
                        "state": snap,
                    })
                else:
                    await websocket.send_json({
                        "type": "error",
                        "message": "No game state available yet",
                    })

    except WebSocketDisconnect:
        pass
    except Exception as e:
        log.warning("WS error game=%s client=%s: %s", game_id, client_id, e)
    finally:
        await game_manager.disconnect(game_id, client_id)


# ────────────────────────────────────────────────────────────
# REST endpoints  —  called by Forge IPC / policy server
# ────────────────────────────────────────────────────────────

class PushStateRequest(BaseModel):
    """Game state snapshot pushed from Forge."""
    game_id: str = ""
    turn: int = 0
    phase: str = ""
    active_player: int = 0
    players: list = Field(default_factory=list)
    stack: list = Field(default_factory=list)
    legal_actions: list = Field(default_factory=list)


class DecisionEvent(BaseModel):
    """AI decision event pushed after policy inference."""
    action: str
    action_index: int = 0
    confidence: float = 0.0
    probabilities: dict = Field(default_factory=dict)
    inference_ms: float = 0.0


class GameOverEvent(BaseModel):
    """Game-over event."""
    winner_seat: int = -1
    reason: str = ""
    turns_played: int = 0


@router.post("/api/ws/game/{game_id}/push")
async def push_game_state(game_id: str, req: PushStateRequest):
    """Push a game state update to all connected Unity clients."""
    tracker = get_tracker(game_id)
    snapshot = req.dict()
    delta_msg = tracker.update(snapshot)
    sent = await game_manager.broadcast_to_game(game_id, delta_msg)
    return {
        "status": "ok",
        "type": delta_msg["type"],
        "seq": delta_msg["seq"],
        "clients_notified": sent,
        "changes": len(delta_msg.get("changes", [])),
    }


@router.post("/api/ws/game/{game_id}/decision")
async def push_decision(game_id: str, event: DecisionEvent):
    """Push an AI decision event to connected Unity clients."""
    tracker = get_tracker(game_id)
    msg = tracker.make_decision_event(
        action=event.action,
        action_index=event.action_index,
        confidence=event.confidence,
        probabilities=event.probabilities,
        inference_ms=event.inference_ms,
    )
    sent = await game_manager.broadcast_to_game(game_id, msg)
    return {"status": "ok", "clients_notified": sent}


@router.post("/api/ws/game/{game_id}/end")
async def push_game_over(game_id: str, event: GameOverEvent):
    """Signal game over to all connected clients and clean up."""
    tracker = get_tracker(game_id)
    msg = tracker.make_game_over_event(
        winner_seat=event.winner_seat,
        reason=event.reason,
        turns_played=event.turns_played,
    )
    sent = await game_manager.broadcast_to_game(game_id, msg)
    remove_tracker(game_id)
    return {"status": "ok", "clients_notified": sent}


@router.get("/api/ws/stats")
async def ws_stats():
    """Return WebSocket connection statistics."""
    return game_manager.stats()