server/app.py

"""FastAPI server exposing the Tetris OpenEnv environment.

Endpoints:
  POST /reset          — create or reset a session env, return observation + session_id
  POST /step           — apply an action to a session env
  GET  /state          — return a session env's current observation
  GET  /health         — liveness probe
  GET  /manifest       — parsed openenv.yaml as JSON
  POST /run_episode    — run a full episode with the random agent, save replay
  GET  /episodes       — list saved replay metadata sorted by training_step asc
  GET  /episodes/{f}   — return a specific replay JSON
  GET  /training-metrics — GRPO run series from metrics.json (see TRAINING_METRICS_PATH)
  POST /reset_all      — clear all active sessions
"""

from __future__ import annotations

import collections
import json
import os
import random
import sys
import uuid
from typing import Any, Dict, List, Optional

import yaml
from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
_PROJECT_ROOT = os.path.dirname(_THIS_DIR)
if _PROJECT_ROOT not in sys.path:
    sys.path.insert(0, _PROJECT_ROOT)

from environment import EpisodeLogger, TetrisEnv  # noqa: E402
from server.models import (  # noqa: E402
    EpisodeMeta,
    ResetRequest,
    ResetResponse,
    RunEpisodeRequest,
    RunEpisodeResponse,
    StepRequest,
    StepResponse,
)

_SERVE_PORT: int = int(os.getenv("PORT", "7860"))

REPLAYS_DIR: str = os.path.join(_PROJECT_ROOT, "replays")

_TRAINING_METRICS_PATH: str = os.getenv(
    "TRAINING_METRICS_PATH",
    os.path.join(_PROJECT_ROOT, "checkpoints", "metrics.json"),
)

_MANIFEST_PATH = os.path.join(_PROJECT_ROOT, "openenv.yaml")
_manifest_data: Dict[str, Any] | None = None
if os.path.isfile(_MANIFEST_PATH):
    with open(_MANIFEST_PATH, "r", encoding="utf-8") as _mf:
        _manifest_data = yaml.safe_load(_mf)


app = FastAPI(title="Tetris OpenEnv", version="0.1.0")

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


MAX_SESSIONS = 50
_sessions: collections.OrderedDict[str, TetrisEnv] = collections.OrderedDict()


def _get_env(session_id: str) -> TetrisEnv:
    """Retrieve a session env by id, raising 404 if not found."""
    env = _sessions.get(session_id)
    if env is None:
        raise HTTPException(status_code=404, detail=f"Session {session_id!r} not found. Call POST /reset first.")
    _sessions.move_to_end(session_id)
    return env


def _create_session() -> tuple[str, TetrisEnv]:
    """Create a new session, evicting the oldest if at capacity."""
    sid = str(uuid.uuid4())
    if len(_sessions) >= MAX_SESSIONS:
        _sessions.popitem(last=False)
    env = TetrisEnv()
    _sessions[sid] = env
    return sid, env


@app.get("/health")
def health() -> Dict[str, str]:
    """Simple liveness probe."""
    return {"status": "ok", "port": str(_SERVE_PORT), "version": "0.1.0"}


@app.get("/manifest")
def manifest() -> Dict[str, Any]:
    """Return the parsed openenv.yaml as JSON."""
    if _manifest_data is None:
        raise HTTPException(status_code=404, detail="Manifest not found")
    return _manifest_data


@app.post("/reset", response_model=ResetResponse)
def reset(req: ResetRequest | None = None) -> ResetResponse:
    """Create a new session or reset an existing one. Returns session_id + observation."""
    sid: str | None = req.session_id if req else None

    if sid and sid in _sessions:
        env = _sessions[sid]
        obs = env.reset()
        _sessions.move_to_end(sid)
    else:
        sid, env = _create_session()
        obs = env.reset()

    return ResetResponse(session_id=sid, observation=obs)


@app.post("/step", response_model=StepResponse)
def step(req: StepRequest) -> StepResponse:
    """Apply a single action to a session environment."""
    env = _get_env(req.session_id)
    try:
        obs, reward, done, info = env.step(req.action)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    return StepResponse(observation=obs, reward=reward, done=done, info=info)


@app.get("/state")
def state(session_id: Optional[str] = Query(None)) -> Dict[str, Any]:
    """Return the current observation of a session environment."""
    if session_id is None:
        if not _sessions:
            raise HTTPException(status_code=404, detail="No active sessions. Call POST /reset first.")
        sid = next(reversed(_sessions))
    else:
        sid = session_id
    env = _get_env(sid)
    return env.get_observation()


@app.post("/reset_all")
def reset_all() -> Dict[str, str]:
    """Clear all active sessions."""
    count = len(_sessions)
    _sessions.clear()
    return {"status": "ok", "sessions_cleared": str(count)}


def _random_agent(_obs: Dict[str, Any]) -> str:
    """Built-in baseline agent that samples a uniformly random action."""
    return random.choice(["move_left", "move_right", "rotate", "place"])


@app.post("/run_episode", response_model=RunEpisodeResponse)
def run_episode(req: RunEpisodeRequest | None = None) -> RunEpisodeResponse:
    """Run one full episode with the built-in random agent on a fresh env."""
    training_step: int = int(req.training_step) if (req and req.training_step is not None) else 0
    agent_label: str = str(req.agent_label) if (req and req.agent_label is not None) else "random"

    local_env = TetrisEnv()
    logger = EpisodeLogger()
    logger.reset()
    obs = local_env.reset()
    total_reward: float = 0.0
    total_lines: int = 0

    safety_cap = 5000
    steps_taken = 0
    placement_turn = 0
    control_actions_taken = 0
    while not local_env.done and steps_taken < safety_cap:
        action = _random_agent(obs)
        board_before = local_env.board.to_list() if action == "place" else []
        piece_info = obs["current_piece"] if action == "place" else {}
        obs, reward, done, _info = local_env.step(action)
        if action == "place":
            placement_turn += 1
            board_after = local_env.board.to_list()
            lines_this_step = int(_info.get("lines_cleared_this_step", 0))
            total_lines += lines_this_step
            logger.log_step(
                turn=placement_turn,
                action="place",
                control_actions_taken=control_actions_taken,
                board_before=board_before,
                piece_info=piece_info,
                board_after=board_after,
                reward=reward,
                lines_cleared=lines_this_step,
                score=local_env.score,
                done=done,
            )
            control_actions_taken = 0
        else:
            control_actions_taken += 1
        total_reward += float(reward["total"])
        steps_taken += 1

    replay_path = logger.save(
        local_env.episode_id,
        training_step,
        agent_label=agent_label,
        output_dir=REPLAYS_DIR,
    )
    return RunEpisodeResponse(
        episode_id=local_env.episode_id,
        training_step=training_step,
        agent_label=agent_label,
        total_reward=float(total_reward),
        total_turns=int(placement_turn),
        total_lines=int(total_lines),
        replay_file=os.path.basename(replay_path),
    )


@app.get("/episodes", response_model=List[EpisodeMeta])
def list_episodes() -> List[EpisodeMeta]:
    """List metadata for all saved replays, sorted by training_step ascending."""
    if not os.path.isdir(REPLAYS_DIR):
        return []
    metas: List[EpisodeMeta] = []
    for name in os.listdir(REPLAYS_DIR):
        if not name.endswith(".json"):
            continue
        path = os.path.join(REPLAYS_DIR, name)
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
            metas.append(
                EpisodeMeta(
                    filename=name,
                    episode_id=str(data.get("episode_id", "")),
                    training_step=int(data.get("training_step", 0)),
                    agent_label=str(data.get("agent_label", "random")),
                    total_turns=int(data.get("total_turns", 0)),
                    total_score=int(data.get("total_score", 0)),
                    total_lines=int(data.get("total_lines", 0)),
                )
            )
        except (OSError, json.JSONDecodeError, ValueError):
            continue
    metas.sort(key=lambda m: m.training_step)
    return metas


@app.get("/training-metrics")
def training_metrics() -> Dict[str, Any]:
    """Return GRPO training curves from `metrics.json` (written by training/grpo_train.py).

    Copy the file from Colab (e.g. `checkpoints/metrics.json`) into the same path in this
    repo, or set env ``TRAINING_METRICS_PATH`` to an absolute path on the server.
    """
    path = _TRAINING_METRICS_PATH
    if not os.path.isfile(path):
        raise HTTPException(
            status_code=404,
            detail=(
                "No training metrics file. Add checkpoints/metrics.json from your training run, "
                "or set TRAINING_METRICS_PATH to the file location."
            ),
        )
    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
    except (OSError, json.JSONDecodeError) as e:
        raise HTTPException(status_code=500, detail=f"Failed to read training metrics: {e}")
    return {
        "steps": data.get("steps", []),
        "mean_rewards": data.get("mean_rewards", []),
        "total_lines": data.get("total_lines", []),
        "total_turns": data.get("total_turns", []),
        "json_valid_rates": data.get("json_valid_rates", []),
    }


@app.get("/episodes/{filename}")
def get_episode(filename: str) -> Dict[str, Any]:
    """Return the full JSON contents of a specific replay file."""
    if "/" in filename or "\\" in filename or filename.startswith(".."):
        raise HTTPException(status_code=400, detail="Invalid filename")
    path = os.path.join(REPLAYS_DIR, filename)
    if not os.path.isfile(path):
        raise HTTPException(status_code=404, detail="Replay not found")
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except (OSError, json.JSONDecodeError) as e:
        raise HTTPException(status_code=500, detail=f"Failed to read replay: {e}")


# ── Serve the built React dashboard if present ────────────────
_DASHBOARD_DIR = os.path.join(_PROJECT_ROOT, "dashboard", "dist")

if os.path.isdir(_DASHBOARD_DIR):
    _DASHBOARD_ASSETS = os.path.join(_DASHBOARD_DIR, "assets")
    if os.path.isdir(_DASHBOARD_ASSETS):
        app.mount("/assets", StaticFiles(directory=_DASHBOARD_ASSETS), name="dashboard-assets")

    @app.get("/")
    def serve_dashboard():
        """Serve the React dashboard at the root URL."""
        return FileResponse(os.path.join(_DASHBOARD_DIR, "index.html"))