"""Data models for the SUPERBEAM OpenEnv environment."""

from __future__ import annotations

from typing import List, Literal, Optional

from openenv.core.env_server import Action, Observation, State
from pydantic import Field

ACTION_NAMES = ("noop", "left", "right", "jump", "right+jump", "left+jump")
ActionName = Literal["noop", "left", "right", "jump", "right+jump", "left+jump"]


class SuperBeamAction(Action):
    """One controller input. It lasts ``frame_skip`` emulator frames (default 4).

    Direction keys are held for the whole action; jump is a tap on the first frame
    (jumping only starts from the ground).
    """

    action: ActionName = Field(description="noop | left | right | jump | right+jump | left+jump")


class SuperBeamObservation(Observation):
    """What the agent sees after reset()/step(). ``reward`` and ``done`` come from the base class."""

    screen: str = Field(default="", description="Labelled ASCII frame, rows 0-25 (P you, E walker, # ground, H pipe (solid), o coin, F flag, C castle, . empty)")
    facts: str = Field(default="", description="One-line factual legend: your columns, ground gaps, walkers, progress")
    frame: List[List[int]] = Field(default_factory=list, description="Raw 32x64 CHIP-8 display bits (row-major, 1 = lit)")
    progress: float = Field(default=0.0, description="Level progress in [0, 1]")
    won: bool = False
    dead: bool = False
    truncated: bool = Field(default=False, description="Episode hit max_frames without winning or dying")
    player_x: int = 0
    player_y: int = 0
    world_x: int = 0
    frames: int = Field(default=0, description="Emulator frames elapsed this episode")
    episode_reward: float = 0.0
    level: str = Field(default="classic", description="Level id: 'classic', an int seed, or 'random'")
    difficulty: Optional[float] = Field(default=None, description="Procedural level difficulty, if supported")
    valid_actions: List[str] = Field(default_factory=lambda: list(ACTION_NAMES))


class SuperBeamState(State):
    """Server-side episode state (GET /state)."""

    level: str = "classic"
    difficulty: Optional[float] = None
    frames: int = 0
    progress: float = 0.0
    won: bool = False
    dead: bool = False
    episode_reward: float = 0.0
    frame_skip: int = 4
    max_frames: int = 600