interface.py

"""Repository-root environment interface for submission validators."""

from __future__ import annotations

from typing import Any, Dict, List, Optional

from env import IndianTrafficEnv
from models import TrafficAction


class TrafficSignalInterface:
    """Thin wrapper exposing reset/step/state around the core environment."""

    def __init__(self, task_id: str = "single_intersection", seed: Optional[int] = 42):
        self._env = IndianTrafficEnv(task_id=task_id)
        self._env.reset(seed=seed, task_id=task_id)

    @staticmethod
    def _dump(model) -> Dict[str, Any]:
        if hasattr(model, "model_dump"):
            return model.model_dump()
        return model.dict()

    def reset(self, seed: Optional[int] = None, task_id: Optional[str] = None) -> Dict[str, Any]:
        return self._dump(self._env.reset(seed=seed, task_id=task_id))

    def step(self, action: str) -> Dict[str, Any]:
        try:
            traffic_action = TrafficAction(action.upper())
        except ValueError:
            traffic_action = TrafficAction.ALL_RED
        state, reward, done, info = self._env.step(traffic_action)
        return {"observation": self._dump(state), "reward": reward, "done": done, "info": info}

    def state(self) -> Dict[str, Any]:
        return self._dump(self._env.get_state())

    def valid_actions(self) -> List[str]:
        return [a.value for a in TrafficAction]


def create_interface(
    task_id: str = "single_intersection", seed: Optional[int] = 42
) -> TrafficSignalInterface:
    """Factory helper for validators expecting a module-level constructor."""
    return TrafficSignalInterface(task_id=task_id, seed=seed)