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

from .models import WordGameAction, WordGameObservation, WordGameState


class WordGameEnv(EnvClient[WordGameAction, WordGameObservation, WordGameState]):
    def _step_payload(self, action: WordGameAction) -> dict:
        return {"guess": action.guess}

    def _parse_result(self, payload: dict) -> StepResult[WordGameObservation]:
        obs_data = payload.get("observation", {})
        return StepResult(
            observation=WordGameObservation(
                masked_word=obs_data.get("masked_word", ""),
                guessed_letters=obs_data.get("guessed_letters", []),
                attempts_remaining=obs_data.get("attempts_remaining", 6),
                message=obs_data.get("message", ""),
                done=payload.get("done", False),
                reward=payload.get("reward"),
            ),
            reward=payload.get("reward"),
            done=payload.get("done", False),
        )

    def _parse_state(self, payload: dict) -> WordGameState:
        return WordGameState(
            episode_id=payload.get("episode_id"),
            step_count=payload.get("step_count", 0),
            target_word=payload.get("target_word", ""),
            max_attempts=payload.get("max_attempts", 6),
        )