src/environment.py

"""
Email Triage OpenEnv Environment.

Implements a realistic email management simulation where AI agents learn to
triage emails into appropriate folders.  Supports deterministic reset with
seed, multi-step episodes, and multi-component reward shaping.

Usage:
    >>> env = EmailTriageEnv(task_id="basic_triage", seed=42)
    >>> obs = env.reset()
    >>> action = Action(action_type="move", email_id=0, target_folder="work")
    >>> obs, reward, done, info = env.step(action)
"""

from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional, Tuple

from src.data_generator import generate_realistic_emails
from src.models import (
    AVAILABLE_FOLDERS,
    MAX_STEPS_PER_EPISODE,
    Action,
    Email,
    Observation,
    Reward,
    StepRecord,
)
from src.reward_shaper import compute_reward

logger = logging.getLogger(__name__)

# ── Task Configuration ───────────────────────────────────────────────────────

TASK_CONFIG: Dict[str, Dict[str, Any]] = {
    "basic_triage": {
        "email_count": 5,
        "difficulty": "easy",
        "max_steps": 15,
        "description": "Sort 5 emails into work vs. spam",
    },
    "multi_folder_triage": {
        "email_count": 15,
        "difficulty": "medium",
        "max_steps": 30,
        "description": "Sort 15 emails into 4 folders",
    },
    "advanced_triage_with_urgency": {
        "email_count": 30,
        "difficulty": "hard",
        "max_steps": 50,
        "description": "Sort 30 emails with VIP and urgency handling",
    },
}


class EmailTriageEnv:
    """Email triage environment following the OpenEnv interface.

    The agent observes an inbox of emails and must sort them into the correct
    folders.  Each action processes one email.  The episode ends when the inbox
    is empty or the step limit is reached.

    Attributes:
        task_id: Which task configuration to use.
        seed: Random seed for deterministic episode generation.
        current_step: Steps taken so far in this episode.
        max_steps: Maximum steps allowed for this task.
        done: Whether the current episode has ended.
        history: List of StepRecords for grading.
    """

    def __init__(
        self,
        task_id: str = "basic_triage",
        seed: Optional[int] = None,
    ) -> None:
        if task_id not in TASK_CONFIG:
            raise ValueError(
                f"Unknown task_id '{task_id}'. "
                f"Available: {list(TASK_CONFIG.keys())}"
            )

        self.task_id = task_id
        self.seed = seed
        self._config = TASK_CONFIG[task_id]

        # Episode state — populated on reset()
        self._inbox: List[Email] = []
        self._folders: Dict[str, List[Email]] = {f: [] for f in AVAILABLE_FOLDERS}
        self._flagged_emails: set = set()
        self._action_history: List[Action] = []
        self.history: List[StepRecord] = []
        self.current_step: int = 0
        self.max_steps: int = self._config["max_steps"]
        self.done: bool = False
        self._episode_num: int = 0
        self._email_lookup: Dict[int, Email] = {}

    # ── OpenEnv Interface ────────────────────────────────────────────────

    def reset(self, seed: Optional[int] = None) -> Observation:
        """Reset the environment to a fresh episode.

        Args:
            seed: Optional override for the random seed.

        Returns:
            Initial observation with a full inbox.
        """
        if seed is not None:
            self.seed = seed

        self._episode_num += 1
        self.current_step = 0
        self.done = False
        self._action_history = []
        self.history = []
        self._flagged_emails = set()
        self._folders = {f: [] for f in AVAILABLE_FOLDERS}

        # Generate emails deterministically
        effective_seed = (
            self.seed if self.seed is not None
            else self._episode_num * 1000 + hash(self.task_id) % 10000
        )
        self._inbox = generate_realistic_emails(
            count=self._config["email_count"],
            difficulty=self._config["difficulty"],
            seed=effective_seed,
        )
        self._email_lookup = {e.id: e for e in self._inbox}

        logger.info(
            "Reset environment: task=%s, emails=%d, seed=%s",
            self.task_id, len(self._inbox), effective_seed,
        )
        return self._build_observation()

    def step(self, action: Action) -> Tuple[Observation, Reward, bool, Dict[str, Any]]:
        """Execute one action in the environment.

        Args:
            action: The agent's chosen action.

        Returns:
            Tuple of (observation, reward, done, info).

        Raises:
            RuntimeError: If the episode has already ended.
            ValueError: If the target email doesn't exist in the inbox.
        """
        if self.done:
            raise RuntimeError(
                "Episode is done. Call reset() to start a new episode."
            )

        self.current_step += 1
        info: Dict[str, Any] = {"task_id": self.task_id}

        # Validate email exists
        email = self._email_lookup.get(action.email_id)
        if email is None:
            # Email not found — return zero reward and continue
            reward = Reward(
                value=0.0,
                components={"correctness": 0.0, "efficiency": 0.0},
                reason=f"Email ID {action.email_id} not found in inbox",
            )
            info["error"] = f"Invalid email_id: {action.email_id}"
            self._check_done()
            obs = self._build_observation()
            self._record_step(action, email or Email(
                id=action.email_id, subject="(unknown)",
                sender="unknown@unknown.com",
                sender_domain="unknown.com",
                timestamp="2024-01-01T00:00:00",
            ), reward)
            return obs, reward, self.done, info

        # Compute reward before processing the action
        reward = compute_reward(
            action=action,
            email=email,
            current_step=self.current_step,
            max_steps=self.max_steps,
            action_history=self._action_history,
            flagged_emails=self._flagged_emails,
        )

        # Process the action
        self._process_action(action, email)
        self._action_history.append(action)

        # Record for grading
        self._record_step(action, email, reward)

        # Check termination
        self._check_done()

        obs = self._build_observation()
        info["reward_components"] = reward.components
        info["emails_remaining"] = len(self._inbox)

        return obs, reward, self.done, info

    def state(self) -> Observation:
        """Return the current observation without advancing the environment."""
        return self._build_observation()

    # ── Internal Methods ─────────────────────────────────────────────────

    def _build_observation(self) -> Observation:
        """Construct the current observation from internal state."""
        return Observation(
            inbox_emails=list(self._inbox),
            available_folders=list(AVAILABLE_FOLDERS),
            current_step=self.current_step,
            max_steps=self.max_steps,
            episode_num=self._episode_num,
            done=self.done,
            info={
                "task_id": self.task_id,
                "emails_processed": len(self.history),
                "folders_used": {
                    f: len(emails) for f, emails in self._folders.items() if emails
                },
            },
        )

    def _process_action(self, action: Action, email: Email) -> None:
        """Update internal state based on the agent's action."""
        if action.action_type == "move":
            self._move_email(email, action.target_folder)
        elif action.action_type == "delete":
            self._remove_from_inbox(email)
        elif action.action_type == "mark_spam":
            self._move_email(email, "spam")
        elif action.action_type == "flag":
            self._flagged_emails.add(email.id)
            # Flagging doesn't remove from inbox
        elif action.action_type == "snooze":
            # Snoozing temporarily removes from inbox
            self._remove_from_inbox(email)

    def _move_email(self, email: Email, folder: str) -> None:
        """Move an email from the inbox to a target folder."""
        self._remove_from_inbox(email)
        if folder in self._folders:
            self._folders[folder].append(email)
        else:
            logger.warning("Unknown folder '%s', defaulting to inbox", folder)
            self._folders["inbox"].append(email)

    def _remove_from_inbox(self, email: Email) -> None:
        """Remove an email from the inbox list."""
        self._inbox = [e for e in self._inbox if e.id != email.id]

    def _check_done(self) -> None:
        """Determine whether the episode should end."""
        if self.current_step >= self.max_steps:
            self.done = True
            logger.info("Episode done: max steps (%d) reached", self.max_steps)
        elif not self._inbox:
            self.done = True
            logger.info("Episode done: inbox empty")

    def _record_step(self, action: Action, email: Email, reward: Reward) -> None:
        """Append a step record for later grading."""
        self.history.append(
            StepRecord(
                step_num=self.current_step,
                action=action,
                email=email,
                reward=reward,
                done=self.done,
            )
        )