# 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 Road Traffic Simulator Environment.

The agent controls a single car navigating from a start to end point
in San Francisco, surrounded by background traffic.
"""

from typing import Optional
from pydantic import Field

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


class RoadTrafficSimulatorAction(Action):
    """
    Action for the traffic navigation environment.

    At each intersection, the agent chooses which outgoing road segment to take.
    The index is into the list of available outgoing edges at the current node,
    sorted by edge ID for determinism. Invalid indices are clipped to valid range.
    """

    next_edge_index: int = Field(
        default=0,
        description=(
            "Index (0-based) of the outgoing edge to take at the current intersection. "
            "Must be in [0, available_actions - 1]. Out-of-range values are clipped."
        ),
    )


class RoadTrafficSimulatorObservation(Observation):
    """
    Observation from the traffic navigation environment.

    Includes the agent car's current GPS position, goal position,
    distance to goal, number of route choices available, and a
    bird's-eye-view screenshot of the map showing real-time traffic
    density (green=free-flow → red=congested).
    """

    # Agent position
    lat: float = Field(default=0.0, description="Current latitude of the agent car")
    lon: float = Field(default=0.0, description="Current longitude of the agent car")

    # Goal
    goal_lat: float = Field(default=0.0, description="Destination latitude")
    goal_lon: float = Field(default=0.0, description="Destination longitude")
    distance_to_goal: float = Field(
        default=0.0, description="Straight-line distance to goal in meters"
    )

    # Action space info
    available_actions: int = Field(
        default=1,
        description="Number of valid outgoing edges from current node (action range: [0, available_actions-1])",
    )

    # Visual observation: 400x400 PNG, base64-encoded
    map_screenshot: str = Field(
        default="",
        description=(
            "Base64-encoded PNG (400×400 px) showing the road network coloured by "
            "traffic density. Blue dot = agent, green star = start, red star = goal, "
            "dashed cyan = planned shortest route."
        ),
    )