env/environment.py

"""
Support Ticket Triage — OpenEnv Environment
============================================

Implements the standard OpenEnv interface:
  reset() -> TicketObservation
  step(action: TicketAction) -> (TicketObservation, TicketReward, bool, dict)
  state() -> EnvironmentState
"""

from __future__ import annotations

import copy
import random
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Tuple

from env.data import TICKET_LOOKUP, TICKETS
from env.graders import GRADERS
from env.models import (
    ActionType,
    Department,
    EnvironmentState,
    TicketAction,
    TicketMessage,
    TicketObservation,
    TicketReward,
    UrgencyLevel,
)
from env.tasks import ALL_TASKS, TASK_LOOKUP, TaskSpec

AVAILABLE_ACTIONS = [at.value for at in ActionType]
AVAILABLE_DEPARTMENTS = [d.value for d in Department]
AVAILABLE_URGENCIES = [u.value for u in UrgencyLevel]


class TicketTriageEnv:
    """
    OpenEnv-compliant environment for Support Ticket Triage.

    Args:
        task_name: One of "route", "triage", "resolve". Default: "route".
        ticket_id: Pin to a specific ticket (for reproducibility). 
                   If None, picks randomly from the task's pool.
        seed: Optional RNG seed.
    """

    # ------------------------------------------------------------------
    # Construction
    # ------------------------------------------------------------------

    def __init__(
        self,
        task_name: str = "route",
        ticket_id: Optional[str] = None,
        seed: Optional[int] = None,
    ) -> None:
        if task_name not in TASK_LOOKUP:
            raise ValueError(
                f"Unknown task '{task_name}'. "
                f"Choose from: {list(TASK_LOOKUP.keys())}"
            )
        self._task_spec: TaskSpec = TASK_LOOKUP[task_name]
        self._pinned_ticket_id: Optional[str] = ticket_id
        self._rng = random.Random(seed)

        # Episode state (initialised on reset)
        self._ticket_data: Dict[str, Any] = {}
        self._observation: Optional[TicketObservation] = None
        self._actions_taken: List[Dict[str, Any]] = []
        self._step_number: int = 0
        self._done: bool = False
        self._cumulative_reward: float = 0.0
        self._follow_up_injected: bool = False

    # ------------------------------------------------------------------
    # OpenEnv API
    # ------------------------------------------------------------------

    def reset(self) -> TicketObservation:
        """Reset the environment and return the initial observation."""
        # Pick ticket
        if self._pinned_ticket_id:
            ticket_id = self._pinned_ticket_id
        else:
            ticket_id = self._rng.choice(self._task_spec.ticket_ids)

        self._ticket_data = copy.deepcopy(TICKET_LOOKUP[ticket_id])
        self._actions_taken = []
        self._step_number = 0
        self._done = False
        self._cumulative_reward = 0.0
        self._follow_up_injected = False

        self._observation = TicketObservation(
            ticket_id=ticket_id,
            subject=self._ticket_data["subject"],
            body=self._ticket_data["body"],
            sender_email=self._ticket_data["sender_email"],
            sender_name=self._ticket_data["sender_name"],
            conversation_history=[
                TicketMessage(
                    sender=self._ticket_data["sender_name"],
                    content=self._ticket_data["body"],
                    timestamp=self._now(),
                )
            ],
            current_department=None,
            current_urgency=None,
            tags=[],
            is_escalated=False,
            is_closed=False,
            step_number=0,
            task_name=self._task_spec.name,
            task_description=self._task_spec.description,
            available_actions=AVAILABLE_ACTIONS,
        )
        return copy.deepcopy(self._observation)

    def step(
        self, action: TicketAction
    ) -> Tuple[TicketObservation, TicketReward, bool, Dict[str, Any]]:
        """
        Apply an action and return (observation, reward, done, info).

        Reward is shaped per-step; the terminal reward summarises the episode.
        """
        if self._done:
            raise RuntimeError("Episode is done. Call reset() to start a new one.")
        if self._observation is None:
            raise RuntimeError("Call reset() before step().")

        self._step_number += 1
        self._observation.step_number = self._step_number

        step_reward, info = self._apply_action(action)
        self._actions_taken.append(action.model_dump())

        # Check terminal conditions
        done = self._check_done(action)
        self._done = done

        # On terminal step: compute final episode reward from grader
        if done:
            final_reward = self._run_grader()
            self._cumulative_reward += final_reward.value
            info["final_grader_reward"] = final_reward.model_dump()
            terminal_reward = final_reward
        else:
            # Mid-episode shaped reward
            self._cumulative_reward += step_reward.value
            terminal_reward = step_reward

        # Possibly inject a follow-up message from the customer (hard task)
        self._maybe_inject_follow_up(action)

        return copy.deepcopy(self._observation), terminal_reward, done, info

    def state(self) -> EnvironmentState:
        """Return full internal state (includes hidden ground truth for graders)."""
        if self._observation is None:
            raise RuntimeError("Call reset() before state().")
        return EnvironmentState(
            observation=copy.deepcopy(self._observation),
            ground_truth=self._ticket_data.get("ground_truth", {}),
            cumulative_reward=self._cumulative_reward,
            step_number=self._step_number,
            done=self._done,
            task_name=self._task_spec.name,
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _apply_action(
        self, action: TicketAction
    ) -> Tuple[TicketReward, Dict[str, Any]]:
        """Mutate observation state and return a shaped step reward."""
        info: Dict[str, Any] = {}

        if action.action_type == ActionType.ROUTE:
            if action.department is None:
                return TicketReward(
                    value=0.0,
                    reason="ROUTE requires 'department' field.",
                    partial_scores={},
                ), {"error": "missing department"}
            self._observation.current_department = action.department
            gt_dept = self._ticket_data["ground_truth"]["correct_department"]
            score = 0.3 if action.department == gt_dept else 0.0
            return TicketReward(
                value=score,
                reason=f"Routed to {action.department.value}.",
                partial_scores={"routing": score},
            ), info

        elif action.action_type == ActionType.SET_URGENCY:
            if action.urgency is None:
                return TicketReward(
                    value=0.0,
                    reason="SET_URGENCY requires 'urgency' field.",
                    partial_scores={},
                ), {"error": "missing urgency"}
            self._observation.current_urgency = action.urgency
            gt_urgency = self._ticket_data["ground_truth"]["correct_urgency"]
            score = 0.2 if action.urgency == gt_urgency else 0.05
            return TicketReward(
                value=score,
                reason=f"Set urgency to {action.urgency.value}.",
                partial_scores={"urgency": score},
            ), info

        elif action.action_type == ActionType.TAG:
            if not action.tags:
                return TicketReward(
                    value=0.0,
                    reason="TAG requires non-empty 'tags' list.",
                    partial_scores={},
                ), {"error": "missing tags"}
            for tag in action.tags:
                if tag not in self._observation.tags:
                    self._observation.tags.append(tag)
            required = self._ticket_data["ground_truth"].get("required_tags", set())
            overlap = len(set(self._observation.tags) & required) / max(len(required), 1)
            return TicketReward(
                value=round(0.1 * overlap, 4),
                reason=f"Added tags: {action.tags}. Overlap with required: {overlap:.0%}",
                partial_scores={"tagging": overlap},
            ), info

        elif action.action_type == ActionType.RESPOND:
            if not action.response_text:
                return TicketReward(
                    value=0.0,
                    reason="RESPOND requires 'response_text' field.",
                    partial_scores={},
                ), {"error": "missing response_text"}
            self._observation.conversation_history.append(
                TicketMessage(
                    sender="Support Agent",
                    content=action.response_text,
                    timestamp=self._now(),
                )
            )
            key_topics = self._ticket_data["ground_truth"].get("key_response_topics", set())
            text_lower = action.response_text.lower()
            found = sum(1 for kw in key_topics if kw.lower() in text_lower)
            quality = found / max(len(key_topics), 1)
            return TicketReward(
                value=round(0.15 * quality, 4),
                reason=f"Response addresses {found}/{len(key_topics)} key topics.",
                partial_scores={"response_quality": quality},
            ), info

        elif action.action_type == ActionType.ESCALATE:
            if not action.escalation_reason:
                return TicketReward(
                    value=0.0,
                    reason="ESCALATE requires 'escalation_reason' field.",
                    partial_scores={},
                ), {"error": "missing escalation_reason"}
            self._observation.is_escalated = True
            needs = self._ticket_data["ground_truth"].get("needs_escalation", False)
            score = 0.2 if needs else -0.1  # penalise unnecessary escalation
            return TicketReward(
                value=max(0.0, score),
                reason=(
                    "Escalated correctly." if needs
                    else "Unnecessary escalation (penalised)."
                ),
                partial_scores={"escalation": score},
            ), info

        elif action.action_type == ActionType.CLOSE:
            self._observation.is_closed = True
            note = action.resolution_note or ""
            score = 0.1 if len(note.split()) >= 5 else 0.0
            return TicketReward(
                value=score,
                reason=f"Ticket closed. Resolution note length: {len(note.split())} words.",
                partial_scores={"closure": score},
            ), info

        else:  # NOOP
            return TicketReward(
                value=0.0,
                reason="NOOP action — no state change.",
                partial_scores={},
            ), info

    def _check_done(self, action: TicketAction) -> bool:
        """Episode ends on CLOSE, task-specific completion, or max_steps reached."""
        if action.action_type == ActionType.CLOSE:
            return True
        # Easy task: routing is the only required action — auto-complete after ROUTE
        if self._task_spec.name == "route" and action.action_type == ActionType.ROUTE:
            return True
        if self._step_number >= self._task_spec.max_steps:
            return True
        return False

    def _maybe_inject_follow_up(self, action: TicketAction) -> None:
        """
        For the hard task, inject a customer follow-up after the first RESPOND.
        Simulates a realistic multi-turn support conversation.
        """
        if self._task_spec.name != "resolve":
            return
        if self._follow_up_injected:
            return
        gt = self._ticket_data["ground_truth"]
        follow_up = gt.get("follow_up_message")
        if not follow_up:
            return
        respond_count = sum(
            1 for a in self._actions_taken if a.get("action_type") == ActionType.RESPOND
        )
        if respond_count >= 1:
            self._observation.conversation_history.append(
                TicketMessage(
                    sender=self._observation.sender_name,
                    content=follow_up,
                    timestamp=self._now(),
                )
            )
            self._follow_up_injected = True

    def _run_grader(self) -> TicketReward:
        """Run the task's grader on the completed episode."""
        grader_fn = GRADERS[self._task_spec.grader_name]
        episode = {
            "observation": self._observation,
            "ground_truth": self._ticket_data.get("ground_truth", {}),
            "actions_taken": self._actions_taken,
            "step_number": self._step_number,
        }
        return grader_fn(episode)

    @staticmethod
    def _now() -> str:
        return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")

    # ------------------------------------------------------------------
    # Class-level metadata
    # ------------------------------------------------------------------

    @classmethod
    def list_tasks(cls) -> List[Dict[str, str]]:
        return [
            {
                "name": t.name,
                "display_name": t.display_name,
                "difficulty": t.difficulty,
                "max_steps": str(t.max_steps),
            }
            for t in ALL_TASKS
        ]