server/environment.py

"""Main BugTriageEnv implementation following the OpenEnv API shape."""

from __future__ import annotations

from typing import Any, Optional

from models import (
    ActionModel,
    CurrentTicketModel,
    ObservationModel,
    QueueStatsModel,
    RewardModel,
    StateModel,
    TicketStateModel,
)
from server.tasks import TaskDefinition, load_task

DEFAULT_SEED = 42
DEFAULT_TASK_ID = "bug_triage_easy"
CRITICAL_SEVERITIES = {"sev0", "sev1"}


class RewardCalculator:
    """Calculates rewards for bug triage actions."""

    BASE_REWARD = 0.50

    CORRECT_SEVERITY = 0.20
    CORRECT_PRIORITY = 0.15
    CORRECT_COMPONENT = 0.15
    CORRECT_TEAM = 0.10
    CORRECT_DUPLICATE = 0.15
    CORRECT_REQUEST_INFO = 0.10
    CORRECT_ESCALATION = 0.15

    INCORRECT_CLOSE_DEFER = 0.20
    MISSED_ESCALATION = 0.15
    INVALID_ACTION = 0.30
    REPEATED_NOOP = 0.10
    UNNECESSARY_SWITCH = 0.05

    ALL_CRITICAL_TRIAGED = 0.10
    BUDGET_EXHAUSTED_CRITICAL_REMAINING = 0.10

    def __init__(self):
        self.action_history = []

    def calculate_step_reward(
        self,
        action: ActionModel,
        ground_truth,
        ticket_id: str,
        is_valid: bool,
        is_critical_ticket: bool,
        is_terminal: bool = False,
        all_critical_triaged: bool = False,
        budget_exhausted_with_critical: bool = False,
    ) -> tuple[float, dict[str, float]]:
        """
        Calculate reward for a single step.
        Rewards are in range [0.0, 1.0].

        Returns:
            tuple of (clamped_reward, breakdown_dict)
        """
        breakdown = {}
        total_reward = self.BASE_REWARD

        if not is_valid:
            breakdown["invalid_action_penalty"] = -self.INVALID_ACTION
            total_reward -= self.INVALID_ACTION
            return self._clamp_reward(total_reward), breakdown

        self.action_history.append((ticket_id, action.action_type))

        if action.action_type == "classify" and action.classify:
            if action.classify.severity == ground_truth.true_severity:
                breakdown["correct_severity"] = self.CORRECT_SEVERITY
                total_reward += self.CORRECT_SEVERITY

            if action.classify.priority == ground_truth.true_priority:
                breakdown["correct_priority"] = self.CORRECT_PRIORITY
                total_reward += self.CORRECT_PRIORITY

            if action.classify.component == ground_truth.true_component:
                breakdown["correct_component"] = self.CORRECT_COMPONENT
                total_reward += self.CORRECT_COMPONENT

        if action.action_type == "assign" and action.assign:
            if action.assign.team == ground_truth.true_assignee_team:
                breakdown["correct_team"] = self.CORRECT_TEAM
                total_reward += self.CORRECT_TEAM

        if action.action_type == "mark_duplicate" and action.mark_duplicate:
            if ground_truth.duplicate_of == action.mark_duplicate.canonical_ticket_id:
                breakdown["correct_duplicate"] = self.CORRECT_DUPLICATE
                total_reward += self.CORRECT_DUPLICATE

        if action.action_type == "request_info" and action.request_info:
            if ground_truth.needs_more_info:
                breakdown["correct_request_info"] = self.CORRECT_REQUEST_INFO
                total_reward += self.CORRECT_REQUEST_INFO

        if action.action_type == "escalate_incident" and action.escalate_incident:
            if ground_truth.true_severity in ["sev0", "sev1"]:
                breakdown["correct_escalation"] = self.CORRECT_ESCALATION
                total_reward += self.CORRECT_ESCALATION

        if action.action_type in ["close", "defer"]:
            if not ground_truth.duplicate_of and ground_truth.needs_more_info:
                breakdown["incorrect_close_defer_penalty"] = -self.INCORRECT_CLOSE_DEFER
                total_reward -= self.INCORRECT_CLOSE_DEFER

        if is_critical_ticket and action.action_type != "escalate_incident":
            if ground_truth.true_severity in ["sev0", "sev1"]:
                breakdown["missed_escalation_penalty"] = -self.MISSED_ESCALATION
                total_reward -= self.MISSED_ESCALATION

        if len(self.action_history) >= 3:
            last_three = self.action_history[-3:]
            if last_three[0] == last_three[1] == last_three[2]:
                breakdown["repeated_noop_penalty"] = -self.REPEATED_NOOP
                total_reward -= self.REPEATED_NOOP

        if is_terminal:
            if all_critical_triaged:
                breakdown["all_critical_triaged"] = self.ALL_CRITICAL_TRIAGED
                total_reward += self.ALL_CRITICAL_TRIAGED

            if budget_exhausted_with_critical:
                breakdown["budget_exhausted_penalty"] = -self.BUDGET_EXHAUSTED_CRITICAL_REMAINING
                total_reward -= self.BUDGET_EXHAUSTED_CRITICAL_REMAINING

        return self._clamp_reward(total_reward), breakdown

    def _clamp_reward(self, reward: float) -> float:
        """Clamp reward to [0.0, 1.0] range."""
        return max(0.0, min(1.0, reward))

    def reset(self):
        """Reset action history."""
        self.action_history = []


class BugTriageEnv:
    """Environment simulating a software bug triage workflow."""

    def __init__(self) -> None:
        self.current_task: Optional[TaskDefinition] = None
        self.current_ticket_index: int = 0
        self.steps_used: int = 0
        self.episode_done: bool = False
        self.cumulative_reward: float = 0.0
        self.last_action_result: Optional[str] = None
        self.ticket_states: list[dict[str, Any]] = []
        self._all_ticket_ids: set[str] = set()

        self.reward_calculator = RewardCalculator()
        self.metrics: dict[str, int | float] = {}

        self._critical_ticket_ids: set[str] = set()
        self._duplicate_ticket_ids: set[str] = set()
        self._info_needed_ticket_ids: set[str] = set()

        self._severity_scored_tickets: set[str] = set()
        self._critical_severity_scored_tickets: set[str] = set()
        self._priority_scored_tickets: set[str] = set()
        self._component_scored_tickets: set[str] = set()
        self._team_scored_tickets: set[str] = set()
        self._duplicate_scored_tickets: set[str] = set()
        self._info_request_scored_tickets: set[str] = set()
        self._escalated_ticket_ids: set[str] = set()
        self._escalation_scored_tickets: set[str] = set()
        self._sla_met_ticket_ids: set[str] = set()
        self._missed_critical_ticket_ids: set[str] = set()

    def reset(
        self,
        task_id: Optional[str] = None,
        seed: Optional[int] = None,
    ) -> ObservationModel:
        """
        Reset the environment to start a new episode.

        Args:
            task_id: One of `bug_triage_easy`, `bug_triage_medium`, `bug_triage_hard`.
            seed: Deterministic shuffle seed.
        """
        resolved_task_id = task_id or DEFAULT_TASK_ID
        resolved_seed = DEFAULT_SEED if seed is None else seed

        self.current_task = load_task(resolved_task_id, seed=resolved_seed)
        self._reset_episode_state()
        self._index_ticket_groups()
        self._reset_dedup_metric_sets()
        self.metrics = self._build_metrics()

        return self._get_observation()

    def step(
        self,
        action: ActionModel,
    ) -> tuple[ObservationModel, RewardModel, bool, dict[str, Any]]:
        """Execute one step and return `(observation, reward, done, info)`."""
        if self.episode_done:
            raise RuntimeError("Episode is done. Call reset() to start new episode.")
        if self.current_task is None:
            raise RuntimeError("No task loaded. Call reset() first.")

        current_ticket = self.current_task.tickets[self.current_ticket_index]
        ground_truth = self._ground_truth_or_raise(current_ticket.ticket_id)
        is_valid, validation_error = self._validate_action(action)
        is_critical = self._is_critical_severity(ground_truth.true_severity)

        step_reward, breakdown = self.reward_calculator.calculate_step_reward(
            action=action,
            ground_truth=ground_truth,
            ticket_id=current_ticket.ticket_id,
            is_valid=is_valid,
            is_critical_ticket=is_critical,
        )

        self.cumulative_reward += step_reward
        self.steps_used += 1
        self.metrics["steps_used"] = self.steps_used

        if is_valid:
            self._update_metrics(action, ground_truth, current_ticket.ticket_id)
        else:
            self.metrics["major_mistakes"] += 1

        self.ticket_states[self.current_ticket_index]["actions_taken"].append(action.action_type)

        if not is_valid:
            self.last_action_result = f"Invalid action: {validation_error}"
        else:
            self.last_action_result = self._execute_action(action, current_ticket.ticket_id)

        self._check_done()
        if self.episode_done:
            terminal_adjustment, terminal_breakdown = self._terminal_step_adjustment()
            if terminal_adjustment:
                step_reward = self.reward_calculator._clamp_reward(step_reward + terminal_adjustment)
                self.cumulative_reward = max(0.0, self.cumulative_reward + terminal_adjustment)
                breakdown.update(terminal_breakdown)

        reward_model = RewardModel(
            step_reward=step_reward,
            cumulative_reward=self.cumulative_reward,
            reward_breakdown=breakdown,
        )

        error_value = validation_error if not is_valid else None
        info = {
            "validation_error": error_value,
            "last_action_error": error_value,
            "metrics": self.metrics.copy(),
        }
        return self._get_observation(), reward_model, self.episode_done, info

    def state(self) -> StateModel:
        """Return the current full environment state."""
        if self.current_task is None:
            raise RuntimeError("No task loaded. Call reset() first.")

        ticket_states = [
            TicketStateModel(
                ticket_id=ticket_state["ticket_id"],
                triaged=ticket_state["triaged"],
                actions_taken=ticket_state["actions_taken"],
            )
            for ticket_state in self.ticket_states
        ]

        return StateModel(
            current_task_id=self.current_task.task_id,
            current_ticket_index=self.current_ticket_index,
            total_tickets=len(self.current_task.tickets),
            tickets_state=ticket_states,
            steps_used=self.steps_used,
            steps_remaining=max(0, self.current_task.step_budget - self.steps_used),
            cumulative_reward=self.cumulative_reward,
            episode_done=self.episode_done,
        )

    def close(self) -> None:
        """No-op close hook for runners expecting a closeable environment."""
        self.current_task = None
        self.ticket_states = []
        self._all_ticket_ids = set()
        self.episode_done = True

    def _reset_episode_state(self) -> None:
        if self.current_task is None:
            return

        self.current_ticket_index = 0
        self.steps_used = 0
        self.episode_done = False
        self.cumulative_reward = 0.0
        self.last_action_result = None
        self._all_ticket_ids = {ticket.ticket_id for ticket in self.current_task.tickets}

        self.ticket_states = [
            {"ticket_id": ticket.ticket_id, "triaged": False, "actions_taken": []}
            for ticket in self.current_task.tickets
        ]
        self.reward_calculator.reset()

    def _index_ticket_groups(self) -> None:
        if self.current_task is None:
            return

        self._critical_ticket_ids = {
            ground_truth.ticket_id
            for ground_truth in self.current_task.ground_truths
            if self._is_critical_severity(ground_truth.true_severity)
        }
        self._duplicate_ticket_ids = {
            ground_truth.ticket_id
            for ground_truth in self.current_task.ground_truths
            if ground_truth.duplicate_of is not None
        }
        self._info_needed_ticket_ids = {
            ground_truth.ticket_id
            for ground_truth in self.current_task.ground_truths
            if ground_truth.needs_more_info
        }

    def _reset_dedup_metric_sets(self) -> None:
        self._severity_scored_tickets = set()
        self._critical_severity_scored_tickets = set()
        self._priority_scored_tickets = set()
        self._component_scored_tickets = set()
        self._team_scored_tickets = set()
        self._duplicate_scored_tickets = set()
        self._info_request_scored_tickets = set()
        self._escalated_ticket_ids = set()
        self._escalation_scored_tickets = set()
        self._sla_met_ticket_ids = set()
        self._missed_critical_ticket_ids = set()

    def _build_metrics(self) -> dict[str, int | float]:
        if self.current_task is None:
            return {}

        label_total = sum(
            1 for ground_truth in self.current_task.ground_truths if ground_truth.duplicate_of is None
        )
        critical_total = len(self._critical_ticket_ids)

        return {
            "severity_correct": 0,
            "priority_correct": 0,
            "component_correct": 0,
            "team_correct": 0,
            "duplicate_correct": 0,
            "duplicate_total": 0,
            "duplicate_expected_total": len(self._duplicate_ticket_ids),
            "info_request_correct": 0,
            "info_needed_total": len(self._info_needed_ticket_ids),
            "major_mistakes": 0,
            "incorrect_close_count": 0,
            "critical_severity_correct": 0,
            "critical_severity_total": critical_total,
            "sla_met": 0,
            "sla_total": critical_total,
            "escalation_correct": 0,
            "escalation_total": critical_total,
            "destructive_actions": 0,
            "missed_critical_escalation": 0,
            "steps_used": 0,
            "step_budget": self.current_task.step_budget,
            "label_total": label_total,
            "assignment_total": label_total,
        }

    def _ground_truth_or_raise(self, ticket_id: str):
        if self.current_task is None:
            raise RuntimeError("No task loaded. Call reset() first.")
        ground_truth = self.current_task.get_ground_truth(ticket_id)
        if ground_truth is None:
            raise RuntimeError(f"Ground truth missing for ticket: {ticket_id}")
        return ground_truth

    @staticmethod
    def _is_critical_severity(severity: str) -> bool:
        return severity in CRITICAL_SEVERITIES

    def _get_observation(self) -> ObservationModel:
        """Build an observation from the current internal state."""
        if self.current_task is None or self.episode_done:
            return ObservationModel(
                current_ticket=None,
                queue_stats=QueueStatsModel(
                    remaining_count=0,
                    urgent_count=0,
                    sla_at_risk_count=0,
                ),
                last_action_result=self.last_action_result,
                available_teams=[],
                available_components=[],
                steps_used=self.steps_used,
                steps_remaining=0,
                partial_score=self._calculate_partial_score(),
            )

        current_ticket = self.current_task.tickets[self.current_ticket_index]
        remaining_count = len(self.current_task.tickets) - self.current_ticket_index
        urgent_count = sum(
            1
            for idx in range(self.current_ticket_index, len(self.current_task.tickets))
            if self._is_critical_severity(
                self._ground_truth_or_raise(self.current_task.tickets[idx].ticket_id).true_severity
            )
        )

        return ObservationModel(
            current_ticket=CurrentTicketModel(**current_ticket.model_dump()),
            queue_stats=QueueStatsModel(
                remaining_count=remaining_count,
                urgent_count=urgent_count,
                sla_at_risk_count=urgent_count,
            ),
            last_action_result=self.last_action_result,
            available_teams=self.current_task.available_teams,
            available_components=self.current_task.available_components,
            steps_used=self.steps_used,
            steps_remaining=max(0, self.current_task.step_budget - self.steps_used),
            partial_score=self._calculate_partial_score(),
        )

    def _validate_action(self, action: ActionModel) -> tuple[bool, Optional[str]]:
        """Validate action legality in the current task context."""
        if self.current_task is None:
            return False, "No task loaded"

        if action.action_type == "mark_duplicate" and action.mark_duplicate:
            canonical_id = action.mark_duplicate.canonical_ticket_id
            if canonical_id not in self._all_ticket_ids:
                return False, f"Unknown ticket ID: {canonical_id}"

        if action.action_type == "assign" and action.assign:
            if action.assign.team not in self.current_task.available_teams:
                return False, f"Unknown team: {action.assign.team}"

        if action.action_type == "classify" and action.classify:
            if action.classify.component not in self.current_task.available_components:
                return False, f"Unknown component: {action.classify.component}"

        return True, None

    def _execute_action(self, action: ActionModel, ticket_id: str) -> str:
        """Apply action side-effects and return a short status message."""
        if self.current_task is None:
            return "No task loaded"

        if action.action_type == "next_ticket":
            if self.current_ticket_index < len(self.current_task.tickets) - 1:
                self.current_ticket_index += 1
                return "Moved to next ticket"
            self.current_ticket_index = len(self.current_task.tickets)
            return "No more tickets in queue"

        if action.action_type == "classify" and action.classify:
            self.ticket_states[self.current_ticket_index]["triaged"] = True
            return f"Classified as {action.classify.severity}/{action.classify.priority}"

        if action.action_type == "assign" and action.assign:
            return f"Assigned to team: {action.assign.team}"

        if action.action_type == "mark_duplicate" and action.mark_duplicate:
            self.ticket_states[self.current_ticket_index]["triaged"] = True
            return f"Marked as duplicate of {action.mark_duplicate.canonical_ticket_id}"

        if action.action_type == "request_info" and action.request_info:
            self.ticket_states[self.current_ticket_index]["triaged"] = True
            return f"Requested {action.request_info.info_type}"

        if action.action_type == "defer" and action.defer:
            self.ticket_states[self.current_ticket_index]["triaged"] = True
            return f"Deferred: {action.defer.reason}"

        if action.action_type == "close" and action.close:
            self.ticket_states[self.current_ticket_index]["triaged"] = True
            return f"Closed: {action.close.reason}"

        if action.action_type == "escalate_incident":
            self.ticket_states[self.current_ticket_index]["triaged"] = True
            return "Escalated incident"

        return f"Action executed for {ticket_id}"

    def _update_metrics(self, action: ActionModel, ground_truth, ticket_id: str) -> None:
        """Update deterministic grading metrics for this step."""
        is_duplicate_ticket = ground_truth.duplicate_of is not None
        is_critical_ticket = self._is_critical_severity(ground_truth.true_severity)

        if is_critical_ticket and action.action_type in {"classify", "mark_duplicate", "escalate_incident"}:
            if ticket_id not in self._sla_met_ticket_ids:
                self._sla_met_ticket_ids.add(ticket_id)
                self.metrics["sla_met"] += 1

        if action.action_type == "classify" and action.classify and not is_duplicate_ticket:
            if action.classify.severity == ground_truth.true_severity:
                if ticket_id not in self._severity_scored_tickets:
                    self._severity_scored_tickets.add(ticket_id)
                    self.metrics["severity_correct"] += 1

                if is_critical_ticket and ticket_id not in self._critical_severity_scored_tickets:
                    self._critical_severity_scored_tickets.add(ticket_id)
                    self.metrics["critical_severity_correct"] += 1

            if action.classify.priority == ground_truth.true_priority:
                if ticket_id not in self._priority_scored_tickets:
                    self._priority_scored_tickets.add(ticket_id)
                    self.metrics["priority_correct"] += 1

            if action.classify.component == ground_truth.true_component:
                if ticket_id not in self._component_scored_tickets:
                    self._component_scored_tickets.add(ticket_id)
                    self.metrics["component_correct"] += 1

        if action.action_type == "assign" and action.assign and not is_duplicate_ticket:
            if action.assign.team == ground_truth.true_assignee_team:
                if ticket_id not in self._team_scored_tickets:
                    self._team_scored_tickets.add(ticket_id)
                    self.metrics["team_correct"] += 1

        if action.action_type == "mark_duplicate" and action.mark_duplicate:
            self.metrics["duplicate_total"] += 1
            if (
                ground_truth.duplicate_of == action.mark_duplicate.canonical_ticket_id
                and ticket_id not in self._duplicate_scored_tickets
            ):
                self._duplicate_scored_tickets.add(ticket_id)
                self.metrics["duplicate_correct"] += 1

        if action.action_type == "request_info" and ground_truth.needs_more_info:
            if ticket_id not in self._info_request_scored_tickets:
                self._info_request_scored_tickets.add(ticket_id)
                self.metrics["info_request_correct"] += 1

        if action.action_type in {"close", "defer"}:
            if not ground_truth.duplicate_of and ground_truth.needs_more_info:
                self.metrics["incorrect_close_count"] += 1
                self.metrics["destructive_actions"] += 1

        if action.action_type == "escalate_incident" and is_critical_ticket:
            self._escalated_ticket_ids.add(ticket_id)
            if ticket_id not in self._escalation_scored_tickets:
                self._escalation_scored_tickets.add(ticket_id)
                self.metrics["escalation_correct"] += 1

        if (
            action.action_type == "next_ticket"
            and is_critical_ticket
            and ticket_id not in self._escalated_ticket_ids
            and ticket_id not in self._missed_critical_ticket_ids
        ):
            self._missed_critical_ticket_ids.add(ticket_id)
            self.metrics["missed_critical_escalation"] += 1

    def _check_done(self) -> None:
        if self.current_task is None:
            return
        if self.steps_used >= self.current_task.step_budget:
            self.episode_done = True
        if self.current_ticket_index >= len(self.current_task.tickets):
            self.episode_done = True

    def _terminal_step_adjustment(self) -> tuple[float, dict[str, float]]:
        if self.current_task is None:
            return 0.0, {}

        all_critical_triaged = self._all_critical_triaged()
        budget_exhausted = self.steps_used >= self.current_task.step_budget
        critical_remaining = self._has_critical_remaining()

        if all_critical_triaged:
            bonus = self.reward_calculator.ALL_CRITICAL_TRIAGED
            return bonus, {"all_critical_triaged_bonus": bonus}

        if budget_exhausted and critical_remaining:
            penalty = -self.reward_calculator.BUDGET_EXHAUSTED_CRITICAL_REMAINING
            return penalty, {"budget_exhausted_critical_remaining_penalty": penalty}

        return 0.0, {}

    def _all_critical_triaged(self) -> bool:
        if self.current_task is None:
            return True

        for index, ticket in enumerate(self.current_task.tickets):
            ground_truth = self._ground_truth_or_raise(ticket.ticket_id)
            if self._is_critical_severity(ground_truth.true_severity) and not self.ticket_states[index]["triaged"]:
                return False
        return True

    def _has_critical_remaining(self) -> bool:
        if self.current_task is None:
            return False

        for index in range(self.current_ticket_index, len(self.current_task.tickets)):
            ticket = self.current_task.tickets[index]
            ground_truth = self._ground_truth_or_raise(ticket.ticket_id)
            if self._is_critical_severity(ground_truth.true_severity) and not self.ticket_states[index]["triaged"]:
                return True
        return False

    def _calculate_partial_score(self) -> float:
        """Compute a normalised [0, 1] partial score based on grading metrics so far."""
        if self.steps_used == 0 or self.current_task is None:
            return 0.0

        total_tickets = len(self.current_task.tickets)
        if total_tickets == 0:
            return 0.0

        score = 0.0
        score += self.metrics.get("severity_correct", 0) / total_tickets * 0.25
        score += self.metrics.get("priority_correct", 0) / total_tickets * 0.20
        score += self.metrics.get("component_correct", 0) / total_tickets * 0.25
        score += self.metrics.get("team_correct", 0) / total_tickets * 0.15
        return min(1.0, score)


if __name__ == "__main__":
    env = BugTriageEnv()
    observation = env.reset(task_id=DEFAULT_TASK_ID, seed=DEFAULT_SEED)
    print(
        "Reset complete. Current ticket:",
        observation.current_ticket.ticket_id if observation.current_ticket else "None",
    )
    print(f"Queue stats: {observation.queue_stats}")
    print(f"State: {env.state().current_task_id}, {env.state().total_tickets} tickets")