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

class Email(BaseModel):
    id: int
    sender: str
    subject: str
    body: str
    folder: str  # e.g., "Inbox", "Archive", "Spam", "Work"
    priority: int # 1 (high) to 3 (low)

class CalendarEvent(BaseModel):
    title: str
    start_time: str
    end_time: str

class Observation(BaseModel):
    inbox_count: int
    current_email: Optional[Email] = None
    calendar: List[CalendarEvent] = []
    folders: List[str] = ["Inbox", "Archive", "Spam", "Work", "Social"]

class Action(BaseModel):
    type: str  # "MOVE", "DELETE", "REPLY", "SCHEDULE"
    email_id: int
    target_folder: Optional[str] = None
    reply_text: Optional[str] = None

class Reward(BaseModel):
    value: float
    reason: str