"""Typed OpenEnv client for the Hospital ED environment.

Maintains a persistent WebSocket connection to a running Hospital ED
server and exposes typed ``reset()`` / ``step()`` / ``state()`` methods
that return :class:`HospitalObservation` and :class:`HospitalState`.

Use this from the baseline ``inference.py`` or any other LLM agent
driver, either pointing it at a locally-running server or spinning up
a fresh Docker container via :meth:`HospitalEdEnv.from_docker_image`::

    # Connect to a running server
    with HospitalEdEnv(base_url="http://localhost:8000") as client:
        result = client.reset(task="surge", seed=0)
        while not result.done:
            result = client.step(HospitalAction(action=0))

    # Or launch a fresh container (hackathon inference pattern)
    client = HospitalEdEnv.from_docker_image("hospital-ed-env:latest")
"""
from __future__ import annotations

from typing import Any, Dict

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

from models import HospitalAction, HospitalObservation, HospitalState


class HospitalEdEnv(
    EnvClient[HospitalAction, HospitalObservation, HospitalState]
):
    """Client for the Hospital Emergency Department OpenEnv environment.

    This client maintains a persistent WebSocket connection to the
    environment server, enabling efficient multi-step interactions with
    lower latency. Each instance holds a dedicated session on the server.
    """

    # ------------------------------------------------------------------
    # OpenEnv EnvClient hooks
    # ------------------------------------------------------------------
    def _step_payload(self, action: HospitalAction) -> Dict[str, Any]:
        """Serialize ``action`` to the JSON payload expected by ``/step``."""
        return {"action": int(action.action)}

    def _parse_result(
        self, payload: Dict[str, Any]
    ) -> StepResult[HospitalObservation]:
        """Parse a server response into a typed :class:`StepResult`."""
        obs_data = payload.get("observation", payload) or {}

        observation = HospitalObservation(
            done=bool(payload.get("done", obs_data.get("done", False))),
            reward=payload.get("reward", obs_data.get("reward")),
            metadata=obs_data.get("metadata", {}) or {},
            bed_occupancy=list(obs_data.get("bed_occupancy", [])),
            icu_occupancy=list(obs_data.get("icu_occupancy", [])),
            ventilator_status=list(obs_data.get("ventilator_status", [])),
            waiting_queue=[
                list(row) for row in obs_data.get("waiting_queue", [])
            ],
            time_step=int(obs_data.get("time_step", 0) or 0),
            stats=list(obs_data.get("stats", [])),
            action_mask=[bool(x) for x in obs_data.get("action_mask", [])],
            survival_rate=float(obs_data.get("survival_rate", 0.0) or 0.0),
            critical_survival_rate=float(
                obs_data.get("critical_survival_rate", 1.0) or 1.0
            ),
            task_score=float(obs_data.get("task_score", 0.0) or 0.0),
            task=obs_data.get("task"),
        )

        return StepResult(
            observation=observation,
            reward=payload.get("reward"),
            done=bool(payload.get("done", False)),
        )

    def _parse_state(self, payload: Dict[str, Any]) -> HospitalState:
        """Parse a ``/state`` payload into a typed :class:`HospitalState`."""
        return HospitalState(
            episode_id=payload.get("episode_id"),
            step_count=int(payload.get("step_count", 0) or 0),
            total_treated=int(payload.get("total_treated", 0) or 0),
            total_deaths=int(payload.get("total_deaths", 0) or 0),
            total_admitted=int(payload.get("total_admitted", 0) or 0),
            critical_total=int(payload.get("critical_total", 0) or 0),
            critical_saved=int(payload.get("critical_saved", 0) or 0),
            queue_len=int(payload.get("queue_len", 0) or 0),
            general_occupancy=float(payload.get("general_occupancy", 0.0) or 0.0),
            icu_occupancy=float(payload.get("icu_occupancy", 0.0) or 0.0),
            ventilator_utilization=float(
                payload.get("ventilator_utilization", 0.0) or 0.0
            ),
            invalid_action_rate=float(
                payload.get("invalid_action_rate", 0.0) or 0.0
            ),
            episode_reward=float(payload.get("episode_reward", 0.0) or 0.0),
            scenario=payload.get("scenario"),
        )