env.py

"""
env.py — InboxOps OpenEnv Environment.

Implements the standard OpenEnv interface:
  env.reset(scenario_id, task_difficulty) → Observation
  env.step(action)                        → (Observation, Reward, done, info)
  env.state()                             → EpisodeState

Three task difficulties gate what the agent must produce per email:
  EASY   → category only
  MEDIUM → category + priority + routing + suggested_action
  HARD   → all of MEDIUM + draft_body + reply_tone
"""

from __future__ import annotations

import json
import os
import random
from copy import deepcopy
from pathlib import Path
from typing import Any, Optional

from models import (
    Action,
    ActionType,
    CalendarEvent,
    Email,
    EmailCategory,
    EpisodeState,
    Observation,
    Priority,
    Reward,
    SLAPolicy,
    SuggestedAction,
    TaskDifficulty,
)
from graders import (
    GradeResult,
    apply_penalties,
    compute_episode_score,
    grade_task1_action,
    grade_task2_action,
    grade_task3_action,
)


# ─────────────────────────────────────────────
# Data loading
# ─────────────────────────────────────────────

DATA_PATH = Path(__file__).parent / "data" / "inbox_scenarios.json"


def _load_data() -> dict:
    with open(DATA_PATH) as f:
        return json.load(f)


# Step budget per difficulty
STEP_BUDGETS = {
    TaskDifficulty.EASY: 20,
    TaskDifficulty.MEDIUM: 35,
    TaskDifficulty.HARD: 50,
}


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

class InboxOpsEnv:
    """
    InboxOps — OpenEnv-compatible environment for operations inbox management.

    Usage
    -----
    env = InboxOpsEnv()
    obs = env.reset(scenario_id="scenario_001", task_difficulty=TaskDifficulty.HARD)
    while not obs.done:
        action = agent.act(obs)
        obs, reward, done, info = env.step(action)
    summary = env.episode_summary()
    """

    def __init__(self, data_path: Optional[Path] = None):
        self._raw = _load_data() if data_path is None else json.loads(Path(data_path).read_text())
        self._scenarios: dict[str, dict] = {s["id"]: s for s in self._raw["scenarios"]}
        self._ground_truth: dict[str, dict] = self._raw["ground_truth"]

        self._state: Optional[EpisodeState] = None
        self._emails: dict[str, Email] = {}
        self._calendar: list[CalendarEvent] = []
        self._sla_policies: dict[str, dict] = {}
        self._reward_history: list[Reward] = []

    # ── Public API ──────────────────────────────

    def reset(
        self,
        scenario_id: Optional[str] = None,
        task_difficulty: TaskDifficulty = TaskDifficulty.MEDIUM,
        seed: Optional[int] = None,
    ) -> Observation:
        """
        Reset the environment for a new episode.

        Parameters
        ----------
        scenario_id : str | None
            Which scenario to load. If None, picks randomly.
        task_difficulty : TaskDifficulty
            Controls what the agent must produce (EASY / MEDIUM / HARD).
        seed : int | None
            Random seed for reproducibility.

        Returns
        -------
        Observation
            Initial observation containing the full inbox.
        """
        if seed is not None:
            random.seed(seed)

        if scenario_id is None:
            scenario_id = random.choice(list(self._scenarios.keys()))

        scenario = self._scenarios[scenario_id]
        self._load_scenario(scenario)

        self._state = EpisodeState(
            scenario_id=scenario_id,
            task_difficulty=task_difficulty,
            max_steps=STEP_BUDGETS[task_difficulty],
        )
        self._reward_history = []

        return self._build_observation("Episode started. Process all pending emails.")

    def step(self, action: Action) -> tuple[Observation, Reward, bool, dict]:
        """
        Apply an action and advance the environment by one step.

        Returns
        -------
        obs : Observation
        reward : Reward
        done : bool
        info : dict  (grading details, debug info)
        """
        if self._state is None:
            raise RuntimeError("Call reset() before step().")
        if self._state.done:
            raise RuntimeError("Episode is done. Call reset().")

        self._state.step_number += 1
        info: dict[str, Any] = {}

        # ── Validate action ──────────────────────
        validation_error = self._validate_action(action)
        if validation_error:
            reward = Reward(breakdown_notes=validation_error)
            apply_penalties(reward, is_invalid=True)
            self._reward_history.append(reward)
            self._state.rewards_log.append(reward.to_dict())
            obs = self._build_observation(f"INVALID ACTION: {validation_error}")
            self._check_termination()
            return obs, reward, self._state.done, {"error": validation_error}

        # ── Handle control actions ───────────────
        if action.action_type == ActionType.COMPLETE:
            self._state.done = True
            self._state.termination_reason = "agent_completed"
            reward = Reward(breakdown_notes="Agent signalled completion.")
            reward.compute_total()
            obs = self._build_observation("Agent marked episode complete.")
            return obs, reward, True, info

        if action.action_type == ActionType.SKIP:
            reward = Reward(breakdown_notes="Skipped.")
            reward.compute_total()
            obs = self._build_observation(f"Skipped email {action.email_id}.")
            self._check_termination()
            return obs, reward, self._state.done, info

        # ── Grade the action ─────────────────────
        grade_result = self._grade(action)
        reward = grade_result.reward

        # Check for duplicate
        dup_count = self._count_duplicate(action)
        if dup_count > 0:
            apply_penalties(reward, is_duplicate=True, duplicate_action_count=dup_count)

        # ── Update state ─────────────────────────
        self._apply_action_to_state(action)
        self._state.actions_log.append(action.to_dict())
        self._state.rewards_log.append(reward.to_dict())
        self._state.cumulative_score += reward.total
        self._reward_history.append(reward)

        # ── Build observation ────────────────────
        result_msg = self._format_result(action, grade_result)
        info["grade_details"] = grade_result.details
        info["reward_breakdown"] = reward.to_dict()

        self._check_termination()
        obs = self._build_observation(result_msg)
        return obs, reward, self._state.done, info

    def state(self) -> EpisodeState:
        if self._state is None:
            raise RuntimeError("No active episode. Call reset().")
        return deepcopy(self._state)

    def episode_summary(self) -> dict:
        if self._state is None:
            return {}
        processed = len(self._state.email_classifications)
        return compute_episode_score(
            rewards=self._reward_history,
            total_emails=len(self._emails),
            emails_processed=processed,
            task_difficulty=self._state.task_difficulty,
        )

    def list_scenarios(self) -> list[str]:
        return list(self._scenarios.keys())

    def get_action_schema(self) -> dict:
        """Return a JSON-serialisable description of the action space."""
        return {
            "action_type": [a.value for a in ActionType],
            "email_id": "string (required for email-targeting actions)",
            "category": [c.value for c in EmailCategory],
            "priority": [p.value for p in Priority],
            "escalation_team": "string (e.g. engineering_oncall, legal_and_founder, bd_team, ...)",
            "suggested_action": [s.value for s in SuggestedAction],
            "draft_body": "string (full reply text, Task 3 only)",
            "reply_tone": "string (e.g. professional_warm, urgent_empathetic, ...)",
            "calendar_slot": "ISO-8601 datetime string (for scheduling actions)",
            "notes": "string (agent reasoning, not graded)",
        }

    # ── Private helpers ─────────────────────────

    def _load_scenario(self, scenario: dict) -> None:
        self._emails = {}
        for e in scenario["emails"]:
            email = Email(
                id=e["id"],
                from_address=e["from"],
                subject=e["subject"],
                body=e["body"],
                timestamp=e["timestamp"],
                thread_id=e["thread_id"],
                has_attachment=e.get("has_attachment", False),
                labels=e.get("labels", []),
            )
            self._emails[email.id] = email

        self._calendar = [
            CalendarEvent(
                id=c["id"],
                title=c["title"],
                start=c["start"],
                end=c["end"],
                attendees=c["attendees"],
            )
            for c in scenario.get("calendar_events", [])
        ]

        self._sla_policies = scenario.get("sla_policies", {})

    def _validate_action(self, action: Action) -> Optional[str]:
        """Return an error message if the action is invalid, else None."""
        if action.action_type in (
            ActionType.CLASSIFY,
            ActionType.SET_PRIORITY,
            ActionType.ROUTE,
            ActionType.SET_ACTION,
            ActionType.DRAFT_REPLY,
            ActionType.ARCHIVE,
            ActionType.ESCALATE,
            ActionType.SKIP,
        ):
            if not action.email_id:
                return "email_id required for this action type."
            if action.email_id not in self._emails:
                return f"Unknown email_id: {action.email_id}"

        if action.action_type == ActionType.CLASSIFY and action.category is None:
            return "category required for CLASSIFY action."
        if action.action_type == ActionType.SET_PRIORITY and action.priority is None:
            return "priority required for SET_PRIORITY action."
        if action.action_type == ActionType.DRAFT_REPLY and not action.draft_body:
            return "draft_body required for DRAFT_REPLY action."

        return None

    def _grade(self, action: Action) -> GradeResult:
        """Route to the correct task grader."""
        email_id = action.email_id
        gt = self._ground_truth.get(email_id, {})

        if not gt:
            # No ground truth → neutral reward
            r = Reward(breakdown_notes=f"No ground truth for {email_id}.")
            r.compute_total()
            return GradeResult(email_id=email_id or "", reward=r, details={})

        difficulty = self._state.task_difficulty
        step = self._state.step_number

        if difficulty == TaskDifficulty.EASY:
            return grade_task1_action(action, gt)
        elif difficulty == TaskDifficulty.MEDIUM:
            return grade_task2_action(action, gt, step)
        else:
            return grade_task3_action(action, gt, step)

    def _apply_action_to_state(self, action: Action) -> None:
        eid = action.email_id
        if not eid:
            return
        if action.category:
            self._state.email_classifications[eid] = action.category
        if action.priority:
            self._state.email_priorities[eid] = action.priority
        if action.escalation_team:
            self._state.email_routings[eid] = action.escalation_team
        if action.suggested_action:
            self._state.email_actions[eid] = action.suggested_action
        if action.draft_body:
            self._state.email_drafts[eid] = action.draft_body
        if action.calendar_slot:
            self._state.booked_slots.append(action.calendar_slot)

    def _count_duplicate(self, action: Action) -> int:
        """Count how many times this exact (email_id, action_type) pair appeared before."""
        eid = action.email_id
        atype = action.action_type.value
        return sum(
            1
            for prev in self._state.actions_log
            if prev.get("email_id") == eid and prev.get("action_type") == atype
        )

    def _check_termination(self) -> None:
        if self._state.done:
            return
        remaining = self._state.max_steps - self._state.step_number
        if remaining <= 0:
            self._state.done = True
            self._state.termination_reason = "step_budget_exhausted"
            return
        # Auto-complete if all emails processed
        processed = set(self._state.email_classifications.keys())
        if processed >= set(self._emails.keys()):
            self._state.done = True
            self._state.termination_reason = "all_emails_processed"

    def _build_observation(self, last_action_result: str) -> Observation:
        processed = set(self._state.email_classifications.keys())
        pending = [eid for eid in self._emails if eid not in processed]

        n_steps = self._state.max_steps - self._state.step_number
        running_score = (
            self._state.cumulative_score / max(self._state.step_number, 1)
        )

        return Observation(
            step_number=self._state.step_number,
            task_difficulty=self._state.task_difficulty,
            emails=list(self._emails.values()),
            pending_email_ids=pending,
            processed_email_ids=list(processed),
            calendar_events=self._calendar,
            active_sla_policies=self._sla_policies,
            last_action_result=last_action_result,
            step_budget_remaining=n_steps,
            score_so_far=min(1.0, max(0.0, running_score)),
            done=self._state.done,
        )

    def _format_result(self, action: Action, result: GradeResult) -> str:
        lines = [
            f"✓ Action: {action.action_type.value} on {action.email_id}",
            f"  Score this step: {result.reward.total:.3f}",
            f"  {result.reward.breakdown_notes}",
        ]
        return "\n".join(lines)