"""Typed OpenEnv wire models for MATH."""

from __future__ import annotations

from typing import Any

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

from math_env.interface import environment_description


class MathAction(Action):
    """The single final-answer string submitted for an episode."""

    answer: str = Field(
        min_length=1,
        description=(
            "One complete final answer string (a number, expression, or set). "
            "No explanation is expected. Plain text or LaTeX is accepted.\n\n"
            + environment_description()
        ),
    )


class MathObservation(Observation):
    """Actor-visible problem or terminal summary without hidden grading evidence.

    Official class and difficulty are post-hoc metadata and are deliberately NOT exposed to
    the solver; they appear only in the evaluator-facing ``MathState``.
    """

    task_id: int
    text: str
    goal: str
    success: bool
    progress: float = Field(ge=0.0, le=1.0)
    terminated: bool
    truncated: bool
    step_limit: int = 1

    @model_validator(mode="before")
    @classmethod
    def _derive_done(cls, values: Any) -> Any:
        if isinstance(values, dict):
            values = dict(values)
            values["done"] = bool(
                values.get("terminated", False) or values.get("truncated", False)
            )
        return values


class MathState(State):
    """Evaluator-facing state for one connection-owned episode.

    Official class and difficulty are retained here as post-hoc metadata for analysis; they
    are never part of the actor-visible observation.
    """

    task_id: int | None = None
    goal: str | None = None
    source_class: str | None = None
    level: str | None = None
    reward: float = 0.0
    success: bool = False
    progress: float = Field(default=0.0, ge=0.0, le=1.0)
    terminated: bool = False
    truncated: bool = False
    step_limit: int = 1
    diagnostics: dict[str, str] = Field(default_factory=dict)
    closed: bool = False