"""Client for the SUPERBEAM OpenEnv environment (WebSocket session)."""

from __future__ import annotations

from typing import Any, Dict

from openenv.core.client_types import StepResult
from openenv.core.env_client import EnvClient

from .models import SuperBeamAction, SuperBeamObservation, SuperBeamState


class SuperBeamEnv(EnvClient[SuperBeamAction, SuperBeamObservation, SuperBeamState]):
    """Typed client.

    Example:
        >>> with SuperBeamEnv(base_url="http://localhost:8000").sync() as env:
        ...     r = env.reset(level="classic")
        ...     r = env.step(SuperBeamAction(action="right"))
        ...     print(r.observation.facts, r.reward, r.done)
    """

    def _step_payload(self, action: SuperBeamAction) -> Dict[str, Any]:
        return {"action": action.action}

    def _parse_result(self, payload: Dict[str, Any]) -> StepResult[SuperBeamObservation]:
        obs_data = dict(payload.get("observation", {}))
        obs_data["done"] = payload.get("done", False)
        obs_data["reward"] = payload.get("reward")
        obs_data["metadata"] = payload.get("metadata") or obs_data.get("metadata", {})
        obs = SuperBeamObservation(**obs_data)
        return StepResult(observation=obs, reward=obs.reward, done=obs.done)

    def _parse_state(self, payload: Dict[str, Any]) -> SuperBeamState:
        return SuperBeamState(**payload)