"""OpenEnv Pydantic models for MiniGridEnv."""

from typing import Any, Optional

from pydantic import BaseModel, ConfigDict, Field

try:
    from openenv.core.env_server.types import Action as _ActionBase
    from openenv.core.env_server.types import Observation as _ObservationBase
    from openenv.core.env_server.types import State as _StateBase
except ImportError:
    _ActionBase = BaseModel
    _ObservationBase = BaseModel
    _StateBase = BaseModel


class MiniGridAction(_ActionBase):
    """Agent action represented as a natural language command."""

    if _ActionBase is BaseModel:
        model_config = ConfigDict(extra="forbid")
        metadata: dict[str, Any] = Field(default_factory=dict)

    command: str = Field(
        ...,
        description="Natural language command (for example: 'go forward', 'turn left', 'pickup').",
    )
    thought: Optional[str] = Field(
        default=None,
        description="Optional reasoning trace for logging/analysis.",
    )


class MiniGridObservation(_ObservationBase):
    """Text observation derived from MiniGrid's egocentric 7x7 view."""

    if _ObservationBase is BaseModel:
        model_config = ConfigDict(extra="forbid")
        done: bool = Field(default=False)
        reward: Optional[float] = Field(default=None)
        metadata: dict[str, Any] = Field(default_factory=dict)

    text: str = Field(..., description="Natural-language environment description.")
    mission: str = Field(..., description="Current mission instruction.")
    step_idx: int = Field(..., ge=0, description="Current 0-indexed step number.")
    steps_remaining: int = Field(..., ge=0, description="Steps left before truncation.")
    max_steps: int = Field(..., ge=1, description="Maximum number of steps in this episode.")
    history: list[dict[str, Any]] = Field(
        default_factory=list,
        description="Recent step history entries for prompt reconstruction.",
    )
    level_name: str = Field(default="", description="Selected BabyAI level short name.")
    last_action: Optional[str] = Field(default=None, description="Canonical last action string.")
    action_success: Optional[bool] = Field(
        default=None,
        description="Whether the last action had an observable effect.",
    )


class MiniGridState(_StateBase):
    """Episode-level state metrics for logging and debugging."""

    if _StateBase is BaseModel:
        model_config = ConfigDict(extra="allow")
        episode_id: Optional[str] = Field(default=None)
        step_count: int = Field(default=0, ge=0)

    level_name: str = Field(default="", description="Selected level name.")
    level_difficulty: int = Field(default=0, ge=0, description="Difficulty stage index.")
    completed: bool = Field(default=False, description="True when mission succeeds.")
    truncated: bool = Field(default=False, description="True when max step budget is exhausted.")
    total_reward: float = Field(default=0.0, description="Cumulative reward this episode.")
    steps_taken: int = Field(default=0, ge=0, description="Total steps executed.")
    optimal_steps: Optional[int] = Field(
        default=None,
        ge=0,
        description="Optional best-path length from the BabyAI bot.",
    )
    efficiency_ratio: Optional[float] = Field(
        default=None,
        description="optimal_steps / steps_taken when available.",
    )
    valid_actions: int = Field(default=0, ge=0, description="Count of valid parsed actions.")
    invalid_actions: int = Field(default=0, ge=0, description="Count of invalid parsed actions.")
    action_distribution: dict[str, int] = Field(
        default_factory=dict,
        description="Histogram of canonical actions used so far.",
    )