env/models.py

from __future__ import annotations
from typing import List, Optional, Literal
from pydantic import BaseModel, Field


class Email(BaseModel):
    id: str
    sender: str
    subject: str
    body: str
    priority: int = Field(..., ge=1, le=5)  # 1=low, 5=critical
    urgency: float = Field(default=1.0, ge=0.0, le=1.0)  # decays each step
    received_at: int = 0  # step number when received
    handled: bool = False
    action_taken: Optional[str] = None


class CalendarEvent(BaseModel):
    time_slot: int  # slot number (0-based)
    title: str
    attendee: Optional[str] = None
    locked: bool = False  # locked = CEO must attend, can't move


class Goal(BaseModel):
    id: str
    description: str
    completed: bool = False
    priority: int = Field(..., ge=1, le=5)
    required_action: str  # "reply", "schedule", "delegate", "ignore"
    target_email_id: str


class Action(BaseModel):
    type: Literal["reply", "schedule", "delegate", "ignore"]
    email_id: str
    note: Optional[str] = None


class StepResult(BaseModel):
    observation: dict
    reward: float
    done: bool
    info: dict


class EnvState(BaseModel):
    current_step: int = 0
    max_steps: int = 6
    current_time: str = "9:00 AM"
    inbox: List[Email] = Field(default_factory=list)
    calendar: List[CalendarEvent] = Field(default_factory=list)
    goals: List[Goal] = Field(default_factory=list)
    history: List[dict] = Field(default_factory=list)
    total_reward: float = 0.0
    done: bool = False