"""Pydantic observation model for ForgeEnv."""
from __future__ import annotations

from typing import Any, Optional

from pydantic import Field

from openenv.core import Observation


class ForgeObservation(Observation):
    """What the agent (or the trainer's rollout function) sees at each step.

    Inherits `done`, `reward`, `metadata` from the OpenEnv `Observation` base.
    """

    current_phase: str = Field(
        ..., description="One of 'drift_gen', 'repair', 'verify', 'done'"
    )
    task_id: str = ""
    task_description: str = ""
    target_category: str = ""
    script_content: str = Field(default="", description="Current state of the script")
    error_trace: Optional[str] = None
    library_versions: dict[str, str] = Field(default_factory=dict)
    reward_breakdown: dict[str, Any] = Field(default_factory=dict)
    held_out_breakdown: dict[str, float] = Field(default_factory=dict)
    episode_step: int = 0
    info: dict[str, Any] = Field(default_factory=dict)