"""
OpenEnv Client for the Shopping Agent Environment.

Connects to the shopping environment server via WebSocket (OpenEnv protocol).
"""

from typing import Any, Dict

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

from openenv_models import ShoppingAction, ShoppingObservation, ShoppingState


class ShoppingEnvClient(EnvClient[ShoppingAction, ShoppingObservation, ShoppingState]):
    """
    WebSocket client for the Shopping Agent Environment.

    Usage (async):
        async with ShoppingEnvClient(base_url="http://localhost:8000") as env:
            result = await env.reset(query="lip balm", product_count=4)
            result = await env.step(ShoppingAction(action_type="view_item", item_ids=["p1"]))

    Usage (sync):
        with ShoppingEnvClient(base_url="http://localhost:8000").sync() as env:
            result = env.reset(query="earbuds")
            result = env.step(ShoppingAction(action_type="buy", item_ids=["p1"]))
    """

    def _step_payload(self, action: ShoppingAction) -> Dict[str, Any]:
        """Convert ShoppingAction to the JSON payload expected by the server."""
        return action.model_dump(exclude_none=True)

    def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ShoppingObservation]:
        """Parse server response into StepResult[ShoppingObservation]."""
        obs_data = payload.get("observation", payload)
        reward = payload.get("reward") or obs_data.get("reward", 0.0)
        done = payload.get("done", obs_data.get("done", False))

        observation = ShoppingObservation(**obs_data)
        return StepResult(
            observation=observation,
            reward=reward,
            done=done,
        )

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