"""
Module: client.py
Purpose: OpenEnv client for connecting to the TriageFlow environment.
Part of: Medical Triage Assistant — OpenEnv Round 1
Author: Team Squirrel

Overview:
    Implements the TriageFlowEnv client that inherits from OpenEnv's EnvClient.
    Handles conversion between typed Python objects and the WebSocket wire format.
    Users import this client to interact with a remote or local TriageFlow server.

Dependencies:
    - openenv.core.env_client: EnvClient base class
    - openenv.core.client_types: StepResult
    - models: TriageAction, TriageObservation, TriageState

Usage:
    from client import TriageFlowEnv
    from models import TriageAction

    async with TriageFlowEnv(base_url="http://localhost:8000") as env:
        result = await env.reset(task_name="basic-triage")
        result = await env.step(TriageAction(action_type="assign_priority", ...))
"""

from openenv.core.env_client import EnvClient
from openenv.core.client_types import StepResult
from models import TriageAction, TriageObservation, TriageState


class TriageFlowEnv(EnvClient[TriageAction, TriageObservation, TriageState]):
    """
    Client for the TriageFlow medical triage environment.

    Provides a type-safe interface for interacting with the TriageFlow
    server via WebSocket. Handles serialization of actions and
    deserialization of observations and state.

    Notes:
        Use .sync() for synchronous access in scripts and notebooks.
    """

    def _step_payload(self, action: TriageAction) -> dict:
        """
        Convert a TriageAction into JSON payload for the server.

        Args:
            action (TriageAction): The typed action object.

        Returns:
            dict: JSON-serializable payload.
        """
        payload = {
            "action_type": action.action_type.value if hasattr(action.action_type, 'value') else str(action.action_type),
            "patient_id": action.patient_id,
        }
        if action.priority_level is not None:
            payload["priority_level"] = action.priority_level.value if hasattr(action.priority_level, 'value') else str(action.priority_level)
        if action.info_field is not None:
            payload["info_field"] = action.info_field.value if hasattr(action.info_field, 'value') else str(action.info_field)
        if action.escalation_reason is not None:
            payload["escalation_reason"] = action.escalation_reason
        return payload

    def _parse_result(self, payload: dict) -> StepResult:
        """
        Parse the server's JSON response into a typed StepResult.

        Args:
            payload (dict): Raw JSON response from the server.

        Returns:
            StepResult: Typed result containing observation, reward, and done flag.
        """
        obs_data = payload.get("observation", payload)
        return StepResult(
            observation=TriageObservation(
                done=payload.get("done", False),
                reward=payload.get("reward"),
                current_patient=obs_data.get("current_patient"),
                queue_length=obs_data.get("queue_length", 0),
                queue_position=obs_data.get("queue_position", 0),
                missing_fields=obs_data.get("missing_fields", []),
                previous_action_feedback=obs_data.get("previous_action_feedback"),
                step_number=obs_data.get("step_number", 0),
                task_name=obs_data.get("task_name", ""),
            ),
            reward=payload.get("reward"),
            done=payload.get("done", False),
        )

    def _parse_state(self, payload: dict) -> TriageState:
        """
        Parse the server's state response into a typed TriageState.

        Args:
            payload (dict): Raw JSON state from the server.

        Returns:
            TriageState: Full internal state object.
        """
        return TriageState(
            episode_id=payload.get("episode_id"),
            step_count=payload.get("step_count", 0),
            task_name=payload.get("task_name", ""),
            patients=payload.get("patients", []),
            assignments=payload.get("assignments", {}),
            escalations=payload.get("escalations", {}),
            info_requests=payload.get("info_requests", []),
            action_history=payload.get("action_history", []),
            current_index=payload.get("current_index", 0),
            max_steps=payload.get("max_steps", 20),
            queue_cleared=payload.get("queue_cleared", False),
        )