client.py

from __future__ import annotations

from typing import Any, Dict

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

from app.openenv_env.models import Action, Observation, EmailEnvState


class EmailAssistantEnvClient(
    EnvClient[Action, Observation, EmailEnvState]
):
    """
    OpenEnv HTTP client for the email assistant environment.

    Handles payload serialization for Action objects and
    parses server responses into strongly-typed local models.
    """

    def _step_payload(self, action: Action) -> Dict[str, Any]:
        """
        Convert Action into JSON-serializable step payload.
        """
        if isinstance(action, dict):
            return dict(action)
        return action.model_dump()

    def _parse_result(self, payload: Dict[str, Any]) -> StepResult[Observation]:
        """
        Parse reset/step server payload into StepResult[Observation].
        """
        # The base reset/step response from the FastAPI OpenEnv router 
        # usually includes observation, reward, done.
        
        obs_payload = payload.get("observation")
        if not obs_payload:
            # Fallback for reset response which might be the observation itself
             observation = Observation.model_validate(payload)
             return StepResult(
                observation=observation,
                reward=None,
                done=False
             )

        observation = Observation.model_validate(obs_payload)

        reward_payload = payload.get("reward")
        reward_value = None
        if isinstance(reward_payload, dict):
            reward_value = reward_payload.get("value")
        elif reward_payload is not None:
            reward_value = reward_payload

        return StepResult(
            observation=observation,
            reward=float(reward_value) if reward_value is not None else None,
            done=bool(payload.get("done", False)),
        )

    def _parse_state(self, payload: Dict[str, Any]) -> EmailEnvState:
        """
        Parse server state payload into EmailEnvState.
        """
        return EmailEnvState.model_validate(payload)