# 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.

"""Supergames Environment Client."""

from typing import Any, Dict

from openenv.core import EnvClient
from openenv.core.client_types import StepResult
from openenv.core.env_server.types import State

from .models import SupergamesAction, SupergamesObservation


class SupergamesEnv(
    EnvClient[SupergamesAction, SupergamesObservation, State]
):
    """
    Client for the Supergames environment (staff allocation across games and sprints).

    Uses a persistent WebSocket session to the environment server. Prefer
    ``SupergamesEnv(base_url=...).sync()`` for a synchronous API in scripts
    and notebooks.

    Example (sync):
        >>> from supergames.models import Assignment, SupergamesAction
        >>> with SupergamesEnv(base_url="http://localhost:8000").sync() as env:
        ...     r = env.reset(task_id=1, seed=42)
        ...     r = env.step(
        ...         SupergamesAction(
        ...             assignments=[Assignment(workItemID="b1", staff=10)],
        ...             reasoning="Triage the chat lag bug first.",
        ...         )
        ...     )

    Example (Docker):
        >>> async def run():
        ...     client = await SupergamesEnv.from_docker_image("supergames-env:latest")
        ...     try:
        ...         await client.connect()
        ...         result = await client.reset(task_id=1, seed=42)
        ...     finally:
        ...         await client.close()
    """

    def _step_payload(self, action: SupergamesAction) -> Dict[str, Any]:
        """Serialize action for the WebSocket ``step`` message."""
        return action.model_dump(mode="json", exclude_none=True)

    def _parse_result(self, payload: Dict[str, Any]) -> StepResult[SupergamesObservation]:
        obs_data: Dict[str, Any] = dict(payload.get("observation", {}))
        obs_data["reward"] = payload.get("reward")
        obs_data["done"] = payload.get("done", False)
        if "metadata" not in obs_data:
            obs_data["metadata"] = {}
        observation = SupergamesObservation.model_validate(obs_data)
        return StepResult(
            observation=observation,
            reward=payload.get("reward"),
            done=payload.get("done", False),
        )

    def _parse_state(self, payload: Dict[str, Any]) -> State:
        return State(
            episode_id=payload.get("episode_id"),
            step_count=payload.get("step_count", 0),
        )