"""Data models for the Executive Assistant Arena Environment."""

from typing import Optional
from pydantic import Field

from openenv.core.env_server.types import Action, Observation, State


class AssistantAction(Action):
    """Action for the assistant environment - tool calls to manage calendar/email."""

    tool: str = Field(
        ...,
        description="Tool to invoke: check_calendar, check_inbox, reschedule, draft_reply, delegate_task, done",
    )
    arguments: dict = Field(
        default_factory=dict,
        description="Tool arguments, e.g. {'event_id': 'mtg_3', 'new_time': '2pm'}",
    )


class AssistantObservation(Observation):
    """Observation from the assistant environment."""

    inbox_summary: str = Field(default="", description="Current emails/messages")
    calendar_view: str = Field(default="", description="Today's schedule as text")
    pending_tasks: list[str] = Field(default_factory=list, description="Unresolved items")
    tool_result: str = Field(default="", description="Output of last tool call")
    conflicts: list[str] = Field(default_factory=list, description="Detected scheduling conflicts")


class AssistantState(State):
    """Internal state tracking for the assistant environment."""

    conflicts_resolved: int = Field(default=0)
    total_conflicts: int = Field(default=0)
    preferences_inferred: int = Field(default=0)
    total_preferences: int = Field(default=0)
    emails_drafted: int = Field(default=0)
    total_emails: int = Field(default=0)
    deadlines_met: int = Field(default=0)
    deadlines_missed: int = Field(default=0)
    unnecessary_actions: int = Field(default=0)
    late_changes_handled: int = Field(default=0)
    total_late_changes: int = Field(default=0)
    cumulative_reward: float = Field(default=0.0)