from pydantic import BaseModel
from enum import Enum
from typing import List, Optional

class ActionType(str, Enum):
    MARK_PRIORITY = "mark_priority"
    MARK_SPAM = "mark_spam"
    AUTO_REPLY = "auto_reply"
    ARCHIVE = "archive"

class Action(BaseModel):
    action_type: ActionType
    email_id: str
    reply_text: Optional[str] = None

class Email(BaseModel):
    id: str
    sender: str
    subject: str
    body: str
    # Making these optional prevents the validation error during observation
    is_spam_truth: Optional[bool] = None 
    is_priority_truth: Optional[bool] = None

class Observation(BaseModel):
    inbox: List[Email]
    steps_taken: int
    max_steps: int