"""
models.py — Typed Pydantic models for Executive Assistant Environment

Defines Action, Observation, and State types used by the OpenEnv spec.
"""

from pydantic import BaseModel
from typing import List, Optional, Dict


# ============================================================
# ACTION — what the agent sends
# ============================================================

class TimeSlot(BaseModel):
    """Proposed meeting time."""
    start_time: str  # ISO format: "2026-04-28T14:00:00"
    end_time: str
    note: Optional[str] = None  # e.g., "This works better for all attendees"


class MeetingDetails(BaseModel):
    """Complete meeting information."""
    participants: List[str]
    start_time: str
    end_time: str
    subject: str
    location: Optional[str] = "Conference Room A"
    proposed_alternatives: Optional[List[TimeSlot]] = None


class AssistantAction(BaseModel):
    """Agent's response to email scenario."""
    email_reply: str  # Draft response to sender
    calendar_action: str  # "book" | "propose_alternatives" | "reschedule" | "decline"
    meeting_details: Optional[MeetingDetails] = None


# ============================================================
# OBSERVATION — what the agent sees
# ============================================================

class Meeting(BaseModel):
    """Existing calendar meeting."""
    id: str
    participants: List[str]
    start_time: str
    end_time: str
    subject: str
    priority: str = "normal"  # "low" | "normal" | "high"


class ContactInfo(BaseModel):
    """Contact metadata."""
    name: str
    email: str
    timezone: str = "America/Los_Angeles"
    title: Optional[str] = None


class EmailInbox(BaseModel):
    """Incoming email request."""
    sender: str
    subject: str
    body: str
    timestamp: str
    priority: str = "normal"


class CalendarState(BaseModel):
    """Current calendar state."""
    existing_meetings: List[Meeting]
    working_hours: Dict[str, str]  # {"monday": "9-17", ...}
    executive_name: str = "Alex Chen"


class AssistantObservation(BaseModel):
    """What the agent receives after reset() or step()."""
    task: Optional[str] = None
    description: Optional[str] = None
    emails: Optional[List[EmailInbox]] = None
    calendar: Optional[CalendarState] = None
    contacts: Optional[Dict[str, ContactInfo]] = None
    action_required: Optional[str] = None
    message: Optional[str] = None


# ============================================================
# STATE — current environment state
# ============================================================

class AssistantState(BaseModel):
    """Current state of the environment."""
    current_task: Optional[str] = None
    emails_pending: int = 0
    episode_done: bool = False
    steps_taken: int = 0
    total_score: float = 0.0