server/models.py

"""Pydantic request/response schemas for the Tetris HTTP API."""

from __future__ import annotations

from typing import Any, Dict, Optional

from pydantic import BaseModel, Field


class ResetRequest(BaseModel):
    """Body for POST /reset."""

    session_id: Optional[str] = Field(None, description="Existing session to reset; omit to create a new one")


class ResetResponse(BaseModel):
    """Response for POST /reset — full observation dict."""

    session_id: str
    observation: Dict[str, Any]


class StepRequest(BaseModel):
    """Body for POST /step."""

    session_id: str = Field(..., description="Session UUID returned by /reset")
    action: str = Field(..., description="One of: move_left, move_right, rotate, place")


class StepResponse(BaseModel):
    """Response for POST /step."""

    observation: Dict[str, Any]
    reward: Dict[str, float]
    done: bool
    info: Dict[str, Any]


class RunEpisodeRequest(BaseModel):
    """Optional body for POST /run_episode."""

    training_step: Optional[int] = 0
    agent_label: Optional[str] = "random"


class RunEpisodeResponse(BaseModel):
    """Response for POST /run_episode."""

    episode_id: str
    training_step: int
    agent_label: str = "random"
    total_reward: float
    total_turns: int
    total_lines: int
    replay_file: str


class EpisodeMeta(BaseModel):
    """Metadata for a single saved replay file."""

    filename: str
    episode_id: str
    training_step: int
    agent_label: str = "random"
    total_turns: int
    total_score: int
    total_lines: int