env.py

"""
CustomerSupportTriageEnv — OpenEnv-compliant environment.

Implements:  reset() · step() · state()
Full episode: agent processes a queue of real-world support tickets,
assigning priority, routing to department, and drafting a reply.
"""
from __future__ import annotations

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

from models import (
    BatchTriageAction, Department, EnvState, Priority,
    ResetRequest, ResetResponse, StepResult, TicketObservation,
    TriageAction, TriageObservation,
)
from tasks import TASKS, Task, build_ticket, grade_action, SupportTicket


class CustomerSupportTriageEnv:
    """
    OpenEnv-compliant environment for Customer Support Triage.

    Episode flow
    ------------
    1. reset(task="easy"|"medium"|"hard") → initial observation (queue of tickets)
    2. step(BatchTriageAction) → grades each action, advances queue, returns reward
       Repeat until done=True
    3. state() → full internal state snapshot

    One episode = one pass through the full ticket queue.
    The agent submits one BatchTriageAction per step covering all remaining tickets,
    OR submits a partial batch (covering ≥1 ticket) and continues.
    """

    ENV_NAME    = "CustomerSupportTriage-v0"
    MAX_STEPS   = 30        # safety cap (normally finishes in 1–3 steps)

    def __init__(self):
        self._task:          Optional[Task]          = None
        self._tickets:       List[SupportTicket]     = []
        self._pending_ids:   List[str]               = []
        self._processed_ids: List[str]               = []
        self._actions_taken: List[Dict[str, Any]]    = []
        self._reward_history: List[float]            = []
        self._episode_reward: float                  = 0.0
        self._step_number:   int                     = 0
        self._done:          bool                    = True
        self._seed:          int                     = 0
        self._ticket_map:    Dict[str, SupportTicket] = {}

    # ── reset ──────────────────────────────────────────────────────────────

    def reset(self, request: Optional[ResetRequest] = None) -> ResetResponse:
        """
        Reset the environment for a new episode.

        Parameters
        ----------
        request.task : "easy" | "medium" | "hard"  (default "easy")
        request.seed : int  (optional; randomises ticket order)
        """
        task_name = "easy"
        seed      = random.randint(0, 99999)

        if request:
            if request.task and request.task in TASKS:
                task_name = request.task
            if request.seed is not None:
                seed = request.seed

        self._seed      = seed
        self._task      = TASKS[task_name]
        self._done      = False
        self._step_number = 0
        self._episode_reward = 0.0
        self._reward_history = []
        self._actions_taken  = []

        # Build tickets, optionally shuffle order
        raw_tickets = [build_ticket(raw) for raw in self._task.tickets]
        rng = random.Random(seed)
        rng.shuffle(raw_tickets)

        self._tickets    = raw_tickets
        self._ticket_map = {t.ticket_id: t for t in raw_tickets}
        self._pending_ids   = [t.ticket_id for t in raw_tickets]
        self._processed_ids = []

        obs = self._make_observation()
        return ResetResponse(observation=obs, task=task_name, seed=seed)

    # ── step ───────────────────────────────────────────────────────────────

    def step(self, action: BatchTriageAction) -> StepResult:
        """
        Process a batch of triage actions.

        The agent must include at least one action for a pending ticket.
        Actions for unknown or already-processed tickets are ignored.

        Returns
        -------
        StepResult with:
          observation : remaining queue (empty when done)
          reward      : mean score across graded actions this step  (0–1)
          done        : True when all tickets are processed
          info        : per-ticket grader breakdowns
        """
        if self._done:
            raise RuntimeError("Episode is done. Call reset() first.")

        self._step_number += 1

        if self._step_number > self.MAX_STEPS:
            self._done = True
            return StepResult(
                observation=self._make_observation(),
                reward=0.0,
                done=True,
                info={"error": "max_steps_exceeded"},
            )

        step_rewards: List[float]      = []
        step_info:    Dict[str, Any]   = {"graded": []}
        newly_processed: List[str]     = []

        for act in action.actions:
            if act.ticket_id not in self._pending_ids:
                # skip duplicate or unknown
                continue

            ticket = self._ticket_map[act.ticket_id]
            reward_info = grade_action(act, ticket)

            step_rewards.append(reward_info.total)
            newly_processed.append(act.ticket_id)

            step_info["graded"].append({
                "ticket_id":        act.ticket_id,
                "priority_score":   reward_info.priority_score,
                "routing_score":    reward_info.routing_score,
                "response_quality": reward_info.response_quality,
                "escalation_score": reward_info.escalation_score,
                "total":            reward_info.total,
                "feedback":         reward_info.feedback,
            })

            self._actions_taken.append({
                "step":         self._step_number,
                "ticket_id":    act.ticket_id,
                "priority":     act.priority.value,
                "department":   act.department.value,
                "needs_human":  act.needs_human,
                "score":        reward_info.total,
            })

        # Update queue
        for tid in newly_processed:
            self._pending_ids.remove(tid)
            self._processed_ids.append(tid)

        # Reward = mean graded score this step (or 0 if nothing was graded)
        step_reward = sum(step_rewards) / len(step_rewards) if step_rewards else 0.0
        self._episode_reward += step_reward
        self._reward_history.append(step_reward)

        self._done = len(self._pending_ids) == 0

        step_info["step_reward"]     = round(step_reward, 4)
        step_info["episode_reward"]  = round(self._episode_reward, 4)
        step_info["pending_count"]   = len(self._pending_ids)
        step_info["processed_count"] = len(self._processed_ids)

        return StepResult(
            observation=self._make_observation(),
            reward=round(step_reward, 4),
            done=self._done,
            info=step_info,
        )

    # ── state ──────────────────────────────────────────────────────────────

    def state(self) -> EnvState:
        """Return a complete snapshot of the current environment state."""
        return EnvState(
            task_name       = self._task.name if self._task else "uninitialized",
            step_number     = self._step_number,
            total_tickets   = len(self._tickets),
            processed       = len(self._processed_ids),
            pending         = len(self._pending_ids),
            episode_reward  = round(self._episode_reward, 4),
            reward_history  = [round(r, 4) for r in self._reward_history],
            actions_taken   = self._actions_taken,
            done            = self._done,
        )

    # ── helpers ─────────────────────────────────────────────────────────────

    def _make_observation(self) -> TriageObservation:
        queue = [
            TicketObservation(
                ticket_id     = t.ticket_id,
                subject       = t.subject,
                body          = t.body,
                customer_name = t.customer_name,
                customer_tier = t.customer_tier,
                created_at    = t.created_at,
                sentiment     = t.sentiment,
                tags          = t.tags,
            )
            for t in self._tickets
            if t.ticket_id in self._pending_ids
        ]
        task_name = self._task.name if self._task else "uninitialized"
        total     = len(self._tickets)
        return TriageObservation(
            queue          = queue,
            processed      = len(self._processed_ids),
            total_tickets  = total,
            task_name      = task_name,
            step_number    = self._step_number,
            time_remaining = max(0, self.MAX_STEPS - self._step_number),
        )