"""Executive Assistant Arena Environment Client."""

from typing import Dict

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

from .models import AssistantAction, AssistantObservation, AssistantState


class ExecAssistantArenaEnv(
    EnvClient[AssistantAction, AssistantObservation, AssistantState]
):
    """
    Client for the Executive Assistant Arena Environment.

    Example:
        >>> with ExecAssistantArenaEnv(base_url="http://localhost:8000") as client:
        ...     result = client.reset(difficulty="medium")
        ...     result = client.step(AssistantAction(tool="check_calendar"))
    """

    def _step_payload(self, action: AssistantAction) -> Dict:
        return {
            "tool": action.tool,
            "arguments": action.arguments,
        }

    def _parse_result(self, payload: Dict) -> StepResult[AssistantObservation]:
        obs_data = payload.get("observation", {})
        observation = AssistantObservation(
            inbox_summary=obs_data.get("inbox_summary", ""),
            calendar_view=obs_data.get("calendar_view", ""),
            pending_tasks=obs_data.get("pending_tasks", []),
            tool_result=obs_data.get("tool_result", ""),
            conflicts=obs_data.get("conflicts", []),
            done=payload.get("done", False),
            reward=payload.get("reward"),
            metadata=obs_data.get("metadata", {}),
        )
        return StepResult(
            observation=observation,
            reward=payload.get("reward"),
            done=payload.get("done", False),
        )

    def _parse_state(self, payload: Dict) -> AssistantState:
        return AssistantState(
            episode_id=payload.get("episode_id"),
            step_count=payload.get("step_count", 0),
            conflicts_resolved=payload.get("conflicts_resolved", 0),
            total_conflicts=payload.get("total_conflicts", 0),
            preferences_inferred=payload.get("preferences_inferred", 0),
            total_preferences=payload.get("total_preferences", 0),
            emails_drafted=payload.get("emails_drafted", 0),
            total_emails=payload.get("total_emails", 0),
            deadlines_met=payload.get("deadlines_met", 0),
            deadlines_missed=payload.get("deadlines_missed", 0),
            unnecessary_actions=payload.get("unnecessary_actions", 0),
            late_changes_handled=payload.get("late_changes_handled", 0),
            total_late_changes=payload.get("total_late_changes", 0),
            cumulative_reward=payload.get("cumulative_reward", 0.0),
        )