env/environment.py

"""
Core environment: ExecOpsEnv
Manages the CEO simulation state machine.
"""

from typing import Optional
from .models import Action, Email, CalendarEvent, Goal, EnvState, StepResult
from .tasks import get_task
from .reward import compute_reward

# 30-minute time slots starting at 9:00 AM
TIME_SLOTS = [
    "9:00 AM", "9:30 AM", "10:00 AM", "10:30 AM",
    "11:00 AM", "11:30 AM", "12:00 PM", "12:30 PM",
    "1:00 PM", "1:30 PM", "2:00 PM", "2:30 PM",
    "3:00 PM", "3:30 PM", "4:00 PM", "4:30 PM",
    "5:00 PM",
]

URGENCY_DECAY = 0.15  # Urgency drops by this amount each step


class ExecOpsEnv:
    def __init__(self, task_id: str = "easy"):
        self.task_id = task_id
        self._task_data = get_task(task_id)
        self._state: Optional[EnvState] = None
        self.reset()

    def reset(self) -> dict:
        """Reinitialize state from task definition. Returns initial observation."""
        task = get_task(self.task_id)

        emails = [Email(**e) for e in task["emails"]]
        calendar = [CalendarEvent(**c) for c in task["calendar"]]
        goals = [Goal(**g) for g in task["goals"]]

        self._state = EnvState(
            current_step=0,
            max_steps=task["max_steps"],
            current_time=TIME_SLOTS[task.get("start_time_slot", 0)],
            inbox=emails,
            calendar=calendar,
            goals=goals,
            history=[],
            total_reward=0.0,
            done=False,
        )
        return self._build_observation()

    def step(self, action: Action) -> dict:
        """
        Execute one step of the environment.
        Order: validate → apply → update state → advance time → decay urgency
               → compute reward → check done → return StepResult
        """
        if self._state is None:
            raise RuntimeError("Environment not initialized. Call reset() first.")
        if self._state.done:
            raise RuntimeError("Environment is done. Call reset() to restart.")

        # 1. Validate action
        valid_action_types = {"reply", "schedule", "delegate", "ignore"}
        if action.type not in valid_action_types:
            raise ValueError(
                f"Invalid action type '{action.type}'. Must be one of: {sorted(valid_action_types)}"
            )
        if not action.email_id:
            raise ValueError("Action is missing required field 'email_id'.")
        email = self._find_email(action.email_id)
        if email is None:
            raise ValueError(f"Email '{action.email_id}' not found in inbox.")
        if email.handled:
            raise ValueError(f"Email '{action.email_id}' is already handled.")

        # 2. Apply action
        email.handled = True
        email.action_taken = action.type

        # Update calendar if scheduling
        if action.type == "schedule":
            self._add_to_calendar(email)

        # 3. Record action in history
        step_entry = {
            "step": self._state.current_step,
            "time": self._state.current_time,
            "action": action.type,
            "email_id": action.email_id,
            "email_subject": email.subject,
            "email_priority": email.priority,
        }

        # 6. Compute reward BEFORE marking goals complete (so reward fn sees pending goals)
        reward = compute_reward(action, email, self._state.goals, self._state)

        # Now mark goals completed
        for goal in self._state.goals:
            if goal.target_email_id == action.email_id and not goal.completed:
                if goal.required_action == action.type:
                    goal.completed = True
        step_entry["reward"] = reward
        self._state.total_reward += reward
        self._state.history.append(step_entry)

        # 4. Advance time
        self._advance_time()

        # 5. Decay urgency on unhandled emails
        self._decay_urgency()

        # 7. Check termination
        self._state.current_step += 1
        self._check_done()

        obs = self._build_observation()
        return {
            "observation": obs,
            "reward": reward,
            "done": self._state.done,
            "info": {
                "step": self._state.current_step,
                "total_reward": self._state.total_reward,
                "action_applied": action.type,
                "email_subject": email.subject,
            },
        }

    def get_state(self) -> dict:
        """Return full serialized state."""
        if self._state is None:
            raise RuntimeError("Environment not initialized.")
        return {
            "current_step": self._state.current_step,
            "max_steps": self._state.max_steps,
            "current_time": self._state.current_time,
            "inbox": [e.model_dump() for e in self._state.inbox],
            "calendar": [c.model_dump() for c in self._state.calendar],
            "goals": [g.model_dump() for g in self._state.goals],
            "history": self._state.history,
            "total_reward": self._state.total_reward,
            "done": self._state.done,
            "task_id": self.task_id,
            "task_name": self._task_data.get("task_name", ""),
            "task_description": self._task_data.get("description", ""),
        }

    # -------------------------
    # Private helpers
    # -------------------------

    def _build_observation(self) -> dict:
        """Build the observation dict that the agent sees."""
        unhandled = [e for e in self._state.inbox if not e.handled]
        pending_goals = [g for g in self._state.goals if not g.completed]
        return {
            "time": self._state.current_time,
            "step": self._state.current_step,
            "max_steps": self._state.max_steps,
            "steps_remaining": self._state.max_steps - self._state.current_step,
            "inbox": [e.model_dump() for e in self._state.inbox],
            "unhandled_count": len(unhandled),
            "calendar": [c.model_dump() for c in self._state.calendar],
            "goals": [g.model_dump() for g in self._state.goals],
            "pending_goals": [g.model_dump() for g in pending_goals],
            "total_reward": self._state.total_reward,
            "done": self._state.done,
            "task_id": self.task_id,
            "task_name": self._task_data.get("task_name", ""),
        }

    def _find_email(self, email_id: str) -> Optional[Email]:
        for email in self._state.inbox:
            if email.id == email_id:
                return email
        return None

    def _advance_time(self):
        """Move to the next 30-minute time slot."""
        try:
            idx = TIME_SLOTS.index(self._state.current_time)
            if idx + 1 < len(TIME_SLOTS):
                self._state.current_time = TIME_SLOTS[idx + 1]
        except ValueError:
            pass

    def _decay_urgency(self):
        """Reduce urgency of all unhandled emails."""
        for email in self._state.inbox:
            if not email.handled:
                email.urgency = max(0.0, email.urgency - URGENCY_DECAY)

    def _check_done(self):
        """Mark done if max steps reached or all goals completed."""
        all_goals_done = all(g.completed for g in self._state.goals)
        if self._state.current_step >= self._state.max_steps or all_goals_done:
            self._state.done = True

    def _add_to_calendar(self, email: Email):
        """Add a scheduling event to the calendar for scheduled emails."""
        # Find next available slot
        locked_slots = {c.time_slot for c in self._state.calendar if c.locked}
        used_slots = {c.time_slot for c in self._state.calendar}
        for slot in range(len(TIME_SLOTS)):
            if slot not in used_slots and slot not in locked_slots:
                self._state.calendar.append(
                    CalendarEvent(
                        time_slot=slot,
                        title=f"Scheduled: {email.subject[:40]}",
                        attendee=email.sender.split("<")[0].strip(),
                        locked=False,
                    )
                )
                break