from __future__ import annotations

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

from freeciv_env.models import FreecivAction, FreecivObservation, FreecivState


class FreecivEnv(EnvClient[FreecivAction, FreecivObservation, FreecivState]):
    def _step_payload(self, action: FreecivAction) -> dict:
        return action.model_dump(exclude_none=True)

    def _parse_result(self, payload: dict) -> StepResult[FreecivObservation]:
        observation = FreecivObservation(**payload["observation"])
        return StepResult(
            observation=observation,
            reward=payload.get("reward"),
            done=payload.get("done", False),
        )

    def _parse_state(self, payload: dict) -> FreecivState:
        return FreecivState(**payload)