# 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 Visual Reasoning Environment."""

from typing import Any, Dict, List, Optional

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


class VisualReasoningAction(Action):
    """Action for the Visual Reasoning environment.

    All fields are Optional[Any] so the server can handle partial/malformed
    JSON from LLMs.  The server is responsible for validation.
    """

    step_type: Optional[Any] = Field(
        default=None,
        description="'advance' to add more content, or 'complete' to finish the explanation",
    )
    narration: Optional[Any] = Field(
        default=None,
        description="Voiceover text describing what this step changes",
    )
    ops: Optional[Any] = Field(
        default=None,
        description=(
            "List of canvas operations. Each op is "
            "{op: str, target_ids: [str], params: {}}"
        ),
    )
    covered_concepts: Optional[Any] = Field(
        default=None,
        description="Items from concept_checklist this step demonstrates",
    )
    intent: Optional[Any] = Field(
        default=None,
        description="Short label describing the agent's intent for this step",
    )


class VisualReasoningObservation(Observation):
    """Observation from the Visual Reasoning environment."""

    # Task info
    task_name: str = Field(default="", description="Difficulty tier: easy, medium, hard, or expert")
    scenario_id: str = Field(default="", description="Identifier for the current scenario")
    goal: str = Field(default="", description="What the explanation should achieve")
    input_data: Dict[str, Any] = Field(default_factory=dict, description="Algorithm input data")
    constraints: List[str] = Field(default_factory=list, description="Task constraints")
    concept_checklist: List[str] = Field(default_factory=list, description="Concepts to cover")
    max_steps: int = Field(default=0, description="Maximum steps allowed")
    # Canvas state
    entities: Dict[str, Dict[str, Any]] = Field(
        default_factory=dict, description="All entities on the canvas keyed by id",
    )
    relations: List[Dict[str, Any]] = Field(
        default_factory=list, description="Edges and relationships between entities",
    )
    layout: Dict[str, Dict[str, float]] = Field(
        default_factory=dict, description="Positions {entity_id: {x, y}}",
    )
    annotations: List[Dict[str, Any]] = Field(
        default_factory=list, description="Annotations attached to entities",
    )
    notes: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="Freestanding notes not tied to entities. "
                    "Each: {note_id: str, text: str, region: Optional[str]}",
    )

    # History
    narration_history: List[str] = Field(
        default_factory=list, description="All narrations so far",
    )
    step_history: List[Dict[str, Any]] = Field(
        default_factory=list, description="Full action history per step",
    )
    concept_coverage: List[str] = Field(
        default_factory=list, description="Concepts covered so far",
    )

    # Step info
    step_id: int = Field(default=0, description="Current step number")
    done_reason: Optional[str] = Field(default=None, description="Why the episode ended")
    remaining_step_budget: int = Field(default=0, description="Steps remaining")

    # Scoring
    score_breakdown: Dict[str, Any] = Field(
        default_factory=dict, description="Per-dimension scores",
    )
    done: bool = Field(default=False, description="Whether the episode is finished")
    reward: float = Field(default=0.0, description="Reward for the last step")
    action_error: Optional[str] = Field(default=None, description="Error from the last action")
    action_penalty: Optional[float] = Field(default=None, description="Penalty applied")