env.py

"""
OpenEnv: Hybrid AI Email & Calendar Assistant Environment
"""

from __future__ import annotations
import random
from datetime import datetime, timedelta
from typing import Optional, Tuple
from pydantic import BaseModel, Field


# ─── Pydantic Models ──────────────────────────────────────────────────────────

class Observation(BaseModel):
    sender: str
    subject: str
    body: str
    priority: str = Field(default="normal", pattern="^(low|normal|high|urgent)$")
    timestamp: str

class Action(BaseModel):
    action_type: str = Field(
        pattern="^(reply|schedule_meeting|archive|flag_urgent|request_info)$"
    )
    content: str = ""
    proposed_time: Optional[str] = None

class Reward(BaseModel):
    score: float = Field(ge=-2.0, le=1.0)
    reason: str


# ─── Sample Emails ────────────────────────────────────────────────────────────

SAMPLE_EMAILS = [
    {
        "sender": "boss@company.com",
        "subject": "URGENT: Submit report by 5 PM TODAY",
        "body": "This is critical. The board needs the Q3 financial report by end of business today. Please drop everything and submit immediately.",
        "priority": "urgent",
        "timestamp": datetime.now().isoformat(),
        "expected_action": "flag_urgent",
    },
    {
        "sender": "colleague@company.com",
        "subject": "Team sync next week?",
        "body": "Hey, would you be available for a 30-min team sync sometime next Tuesday or Wednesday between 10 AM and 2 PM? Let me know your availability.",
        "priority": "normal",
        "timestamp": (datetime.now() - timedelta(hours=2)).isoformat(),
        "expected_action": "schedule_meeting",
    },
    {
        "sender": "noreply@win-prize.biz",
        "subject": "You've WON $1,000,000!!!",
        "body": "CONGRATULATIONS! You have been selected as a winner. Click here to claim your prize. Limited time offer. Act now!!!",
        "priority": "low",
        "timestamp": (datetime.now() - timedelta(hours=5)).isoformat(),
        "expected_action": "archive",
    },
    {
        "sender": "client@partnerco.com",
        "subject": "Regarding our discussion",
        "body": "Hi, following up on our last conversation. Could you clarify the timeline and next steps? Not entirely sure what the plan is from here.",
        "priority": "normal",
        "timestamp": (datetime.now() - timedelta(hours=1)).isoformat(),
        "expected_action": "request_info",
    },
    {
        "sender": "hr@company.com",
        "subject": "Action required: Performance review + schedule 1:1",
        "body": "Your annual performance review is due this Friday. Please complete the self-assessment form AND schedule a 1:1 meeting with your manager before Thursday.",
        "priority": "high",
        "timestamp": (datetime.now() - timedelta(minutes=30)).isoformat(),
        "expected_action": "schedule_meeting",
    },
]


# ─── Calendar ─────────────────────────────────────────────────────────────────

class Calendar:
    def __init__(self):
        self.slots: dict[str, str] = {}  # "HH:MM" → "free" | "busy:<reason>"

    def reset(self):
        self.slots = {f"{h:02d}:00": "free" for h in range(9, 19)}

    def book(self, time_str: str, reason: str = "meeting") -> Tuple[bool, str]:
        key = time_str[:5]
        if key not in self.slots:
            return False, f"Invalid slot: {key}"
        if self.slots[key] != "free":
            return False, f"Slot {key} is already BUSY"
        self.slots[key] = f"busy:{reason}"
        return True, f"Slot {key} booked for {reason}"

    def is_free(self, time_str: str) -> bool:
        return self.slots.get(time_str[:5], "free") == "free"


# ─── Environment ──────────────────────────────────────────────────────────────

class EmailEnv:
    def __init__(self):
        self.calendar = Calendar()
        self._emails = list(SAMPLE_EMAILS)
        self._index = 0
        self._done = False
        self._history: list[dict] = []
        self.calendar.reset()

    def reset(self) -> Observation:
        self._index = 0
        self._done = False
        self._history = []
        self.calendar.reset()
        return self._make_observation()

    def _make_observation(self) -> Observation:
        e = self._emails[self._index]
        return Observation(
            sender=e["sender"],
            subject=e["subject"],
            body=e["body"],
            priority=e["priority"],
            timestamp=e["timestamp"],
        )

    def step(self, action: Action) -> Tuple[Optional[Observation], Reward, bool, dict]:
        expected = self._emails[self._index]["expected_action"]
        reward = self._compute_reward(action, expected)

        # Calendar side-effects
        info = {"calendar_update": None}
        if action.action_type == "schedule_meeting" and action.proposed_time:
            ok, msg = self.calendar.book(action.proposed_time, "meeting")
            info["calendar_update"] = msg

        self._history.append({
            "email": self._emails[self._index]["subject"],
            "action": action.action_type,
            "expected": expected,
            "score": reward.score,
        })

        self._index += 1
        self._done = self._index >= len(self._emails)

        next_obs = None if self._done else self._make_observation()
        return next_obs, reward, self._done, info

    def _compute_reward(self, action: Action, expected: str) -> Reward:
        at = action.action_type
        email = self._emails[self._index]

        if at == expected:
            base = 1.0
            reason = f"Correct action '{at}' for this email."
        elif email["priority"] == "urgent" and at != "flag_urgent":
            return Reward(score=-2.0, reason="MISSED URGENT EMAIL — critical failure.")
        elif at in ("reply", "request_info") and expected in ("reply", "request_info"):
            base = 0.5
            reason = "Partially correct — close but not optimal."
        elif at == "archive" and expected != "archive":
            return Reward(score=-1.0, reason="Wrongly archived a legitimate email.")
        else:
            base = -1.0
            reason = f"Wrong action. Expected '{expected}', got '{at}'."

        # Quality bonus
        if action.content and len(action.content) > 50:
            base = min(1.0, base + 0.5)
            reason += " Good reply quality."
        elif at in ("reply", "request_info") and len(action.content) < 20:
            base = max(-2.0, base - 0.5)
            reason += " Poor reply quality."

        return Reward(score=round(base, 2), reason=reason)

    def state(self) -> dict:
        return {
            "current_index": self._index,
            "total_emails": len(self._emails),
            "done": self._done,
            "history": self._history,
            "calendar": dict(self.calendar.slots),
        }