# client.py — used in training code / run_baseline.py
from typing import Any
from openenv.core.env_client import EnvClient
from openenv.core.client_types import StepResult
from .models import Action, Observation, State


class CodeDebugEnv(EnvClient[Action, Observation, State]):
    """
    Client for the CodeDebug environment.
    Usage:
        async with CodeDebugEnv(base_url="https://your-space.hf.space") as env:
            result = await env.reset(task_id="task_easy")
            result = await env.step(Action(patch="...", task_id="task_easy"))
    """

    def _step_payload(self, action: Action) -> dict[str, Any]:
        return action.model_dump(exclude_none=True)

    def _parse_result(self, payload: dict[str, Any]) -> StepResult[Observation]:
        obs = Observation(**payload.get("observation", payload))
        return StepResult(
            observation=obs,
            reward=payload.get("reward", obs.reward),
            done=payload.get("done", obs.done)
        )

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