"""
OpenEnv Pydantic models for the env/ stack.

Matches the env/environment data shape: observations with prompt, target_name,
step, paper_fold_json; actions as fold dicts with from/to/assignment.
"""
from typing import Optional

from pydantic import ConfigDict, Field

from openenv.core.env_server.types import Action, Observation, State


class OrigamiAction(Action):
    """One fold operation — from_point, to_point, assignment."""

    model_config = ConfigDict(populate_by_name=True)

    from_point: list[float] = Field(
        alias="from",
        description="[x, y] start point of the crease",
    )
    to_point: list[float] = Field(
        alias="to",
        description="[x, y] end point of the crease",
    )
    assignment: str = Field(
        description="'M' (mountain) or 'V' (valley)",
    )


class OrigamiObservation(Observation):
    """Observation from env.environment — prompt, target, step, paper state."""

    prompt: str = Field(default="", description="LLM prompt for the current step")
    target_name: str = Field(default="", description="Name of the target (.fold stem)")
    step: int = Field(default=0, ge=0, description="Current step index")
    paper_fold_json: dict = Field(
        default_factory=dict,
        description="Graph edges (crease pattern state)",
    )


class OrigamiState(State):
    """Server-side episode state."""

    paper: dict = Field(default_factory=dict, description="Paper state")
    target: Optional[str] = Field(default=None, description="Target name")
    step: int = Field(default=0, ge=0, description="Step count")
    mode: str = Field(default="step", description="'step' or 'code_as_policy'")


__all__ = ["OrigamiAction", "OrigamiObservation", "OrigamiState"]