# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
Data models for the Constraint Env Environment.

The constraint_env environment is a simple test environment that echoes back messages.
"""

from openenv.core.env_server.types import Action, Observation, State
from typing import Dict, Any, Optional, Union
import json as _json
from pydantic import field_validator


class ConstraintAction(Action):
    """Output as AST by LLM."""

    ast_output: str

    @field_validator("ast_output", mode="before")
    @classmethod
    def _normalise(cls, v: Any) -> str:
        """Accept dict (from Gradio UI) or str (from LLM/API), always store as JSON string."""
        if isinstance(v, dict):
            return _json.dumps(v)
        if not isinstance(v, str):
            return _json.dumps(v)
        return v


class ConstraintObservation(Observation):
    """Observation from the environment, user prompt and rewards"""

    prompt: str
    done: bool
    reward: float
    info: Dict[str, Any]
    messages: list[Dict[str, Any]] = []


class ConstraintState(State):
    """Current state from the environment"""
    episode_id: Optional[str] = None
    step_count: int = 0
    max_steps: int = 5