"""
Typed data models for the Autonomous Traffic Control Environment.

Uses openenv-core base types (Action, Observation, State) so that the
environment is fully compliant with the OpenEnv specification and can
be consumed by any openenv-core-compatible RL training framework.
"""

from typing import Any, Dict, List, Optional, Union

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


# ---------------------------------------------------------------------------
# Phase and direction constants (exported for external consumers)
# ---------------------------------------------------------------------------

PHASE_NS_GREEN  = 0   # North-South green, East-West red
PHASE_EW_GREEN  = 1   # East-West green, North-South red
PHASE_ALL_RED   = 2   # All approaches red (emergency clearance)
PHASE_NS_YELLOW = 3   # North-South transitioning to red (internal)
PHASE_EW_YELLOW = 4   # East-West transitioning to red (internal)

DIRECTION_NORTH = 0
DIRECTION_SOUTH = 1
DIRECTION_EAST  = 2
DIRECTION_WEST  = 3


# ---------------------------------------------------------------------------
# Action  (extends openenv-core Action)
# ---------------------------------------------------------------------------

class TrafficAction(Action):
    """
    Agent action: set the desired traffic light phase.

    light_phase:
        0 = NS_GREEN  – North + South get green, East + West get red.
        1 = EW_GREEN  – East + West get green, North + South get red.
        2 = ALL_RED   – All approaches red; use for emergency clearance.

    Yellow-light transitions are handled automatically when switching between
    NS_GREEN and EW_GREEN.
    """
    light_phase: int = Field(
        ...,
        ge=0, le=2,
        description="Desired light phase: 0=NS_GREEN | 1=EW_GREEN | 2=ALL_RED",
    )


# ---------------------------------------------------------------------------
# Observation  (extends openenv-core Observation which already has done/reward/metadata)
# ---------------------------------------------------------------------------

class TrafficObservation(Observation):
    """
    Full observation returned by reset() and step().

    Directions index: 0=North, 1=South, 2=East, 3=West

    The `done`, `reward`, and `metadata` fields are inherited from
    openenv.core.env_server.types.Observation.
    """

    # -- Traffic-light state --
    current_phase: int = Field(
        default=0,
        description="Active light phase: 0=NS_GREEN|1=EW_GREEN|2=ALL_RED|3=NS_YELLOW|4=EW_YELLOW",
    )
    time_in_phase: int = Field(
        default=0,
        description="Steps elapsed since last phase change",
    )

    # -- Vehicle queues (one per direction [N, S, E, W]) --
    queue_lengths:     List[int] = Field(default_factory=lambda: [0, 0, 0, 0],
                                         description="Regular vehicle count per approach")
    emergency_queue:   List[int] = Field(default_factory=lambda: [0, 0, 0, 0],
                                         description="Emergency vehicle count per approach")
    emergency_urgency: List[int] = Field(default_factory=lambda: [0, 0, 0, 0],
                                         description="Max urgency (0-10) of waiting emergency vehicles per approach")

    # -- Flow metrics for this step --
    vehicles_passed:  int   = Field(default=0, description="Regular vehicles cleared this step")
    emergency_passed: int   = Field(default=0, description="Emergency vehicles cleared this step")

    # -- Penalty signals --
    total_waiting_time: float = Field(default=0.0, description="Sum of per-vehicle waiting increments this step")
    collision:          bool  = Field(default=False, description="Gridlock-induced collision flag")


# ---------------------------------------------------------------------------
# State  (extends openenv-core State which has episode_id / step_count)
# ---------------------------------------------------------------------------

class TrafficState(State):
    """
    Episode-level cumulative state, returned by state().

    The `episode_id` and `step_count` fields are inherited from
    openenv.core.env_server.types.State.
    """
    task_id: str = Field(default="basic_flow", description="Active task ID")

    # Cumulative episode metrics
    total_vehicles_passed:  int   = Field(default=0)
    total_emergency_passed: int   = Field(default=0)
    total_waiting_time:     float = Field(default=0.0)
    total_emergency_delay:  float = Field(default=0.0, description="Steps emergency vehicles spent waiting")
    total_collisions:       int   = Field(default=0)
    total_phase_changes:    int   = Field(default=0)