environment.py

"""
environment.py — Strict RL-style Customer Support Environment
Handles: step enforcement, repeat detection, fail conditions, reward calculation
"""

from __future__ import annotations
import re
from typing import List, Tuple, Optional

from models import (
    Episode, EpisodeStatus, StepName, StepResult,
    DifficultyLevel, Task
)
from graders.base_grader import BaseGrader


# ── Step order ────────────────────────────────────────────────────────────────
STEP_ORDER = [
    StepName.EMPATHY,
    StepName.COLLECT_INFO,
    StepName.INVESTIGATE,
    StepName.RESOLUTION,
]

# ── Reward constants ───────────────────────────────────────────────────────────
BASE_SCORE_CORRECT   = 1.0
BASE_SCORE_INCORRECT = 0.2
STEP_BONUS           = 0.2          # bonus when step is correct
WRONG_STEP_PENALTY   = 0.3          # wrong action in correct step position
REPEAT_PENALTY       = 0.2          # repeated question / response
SKIP_STEP_PENALTY    = 0.3          # jumped ahead
EARLY_SOLUTION_PENALTY = 0.25       # gave resolution before investigation
EMOTION_IGNORE_PENALTY = 0.25       # angry customer → neutral / cold reply
GENERIC_RESPONSE_PENALTY = 0.15     # "ok", "done", vague one-liners
WRONG_ASSUMPTION_PENALTY = 0.2      # stated wrong facts
LOOP_PENALTY         = 0.35         # agent stuck in loop (same step repeated 2+)

# ── Overlap threshold for repeat detection ─────────────────────────────────────
REPEAT_WORD_OVERLAP_MIN = 10        # words in common → flagged as repeat


class CustomerSupportEnv:
    """
    Strict step-based RL environment for customer support training.
    """

    def __init__(self, task: Task, grader: BaseGrader):
        self.task    = task
        self.grader  = grader
        self.episode = Episode(
            task_id    = task.task_id,
            difficulty = task.difficulty,
        )
        self._response_history: List[str] = []
        self._step_index = 0           # which step we expect next
        self._consecutive_wrong = 0    # wrong attempts at current step

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

    def step(self, agent_response: str) -> Tuple[StepResult, bool]:
        """
        Process one agent response.
        Returns (StepResult, done:bool).
        """
        if self.episode.status != EpisodeStatus.RUNNING:
            raise RuntimeError("Episode is already finished.")

        expected_step = STEP_ORDER[self._step_index]
        result = self._evaluate(agent_response, expected_step)
        self.episode.add_step(result)
        self._response_history.append(agent_response.lower().strip())

        done = False
        if result.correct:
            self._step_index        += 1
            self._consecutive_wrong  = 0
        else:
            self._consecutive_wrong += 1

        # Loop fail: 3 consecutive wrong attempts at the same step
        if self._consecutive_wrong >= 3 and self.episode.status == EpisodeStatus.RUNNING:
            self.episode.status    = EpisodeStatus.FAIL
            self.episode.fail_reason = (
                f"Agent stuck in loop at step '{expected_step.value}' "
                f"({self._consecutive_wrong} consecutive failures)"
            )

        if self.episode.status != EpisodeStatus.RUNNING:
            done = True
        elif self._step_index >= len(STEP_ORDER):
            done = True   # episode.add_step() already set SUCCESS/FAIL

        return result, done

    def reset(self) -> None:
        self.episode = Episode(
            task_id    = self.task.task_id,
            difficulty = self.task.difficulty,
        )
        self._response_history  = []
        self._step_index        = 0
        self._consecutive_wrong = 0

    def summary(self):
        return self.episode.summary()

    # ── Internal evaluation ────────────────────────────────────────────────────

    def _evaluate(self, response: str, expected_step: StepName) -> StepResult:
        penalties: List[str] = []
        total_penalty        = 0.0

        # 1. Grade the response against expected step
        grader_result = self.grader.grade(
            response      = response,
            expected_step = expected_step,
            task          = self.task,
        )
        correct        = grader_result["correct"]
        base_score     = BASE_SCORE_CORRECT if correct else BASE_SCORE_INCORRECT
        detected_action = grader_result.get("detected_action", "unknown")

        # 2. Step bonus
        step_bonus = STEP_BONUS if correct else 0.0

        # 3. Wrong-step penalty
        if not correct:
            total_penalty += WRONG_STEP_PENALTY
            penalties.append(
                f"Wrong action detected ('{detected_action}' "
                f"≠ '{expected_step.value}'): -{WRONG_STEP_PENALTY}"
            )

        # 4. Repeat detection
        if self._is_repeated_response(response):
            total_penalty += REPEAT_PENALTY
            penalties.append(f"Repeated/duplicate response: -{REPEAT_PENALTY}")

        # 5. Early solution penalty
        if self._is_early_solution(response, expected_step):
            total_penalty += EARLY_SOLUTION_PENALTY
            penalties.append(f"Solution given too early: -{EARLY_SOLUTION_PENALTY}")

        # 6. Emotion mismatch penalty
        if self._is_emotion_mismatch(response):
            total_penalty += EMOTION_IGNORE_PENALTY
            penalties.append(
                f"Angry customer ignored (no empathy/de-escalation): "
                f"-{EMOTION_IGNORE_PENALTY}"
            )

        # 7. Generic / too-short response
        if self._is_generic_response(response):
            total_penalty += GENERIC_RESPONSE_PENALTY
            penalties.append(f"Generic/too-short response: -{GENERIC_RESPONSE_PENALTY}")

        # 8. Wrong assumption detection
        wrong_assumption = grader_result.get("wrong_assumption", False)
        if wrong_assumption:
            total_penalty += WRONG_ASSUMPTION_PENALTY
            penalties.append(f"Incorrect assumption stated: -{WRONG_ASSUMPTION_PENALTY}")

        # 9. Skip-step penalty (grader signals this)
        if grader_result.get("skipped_step", False):
            total_penalty += SKIP_STEP_PENALTY
            penalties.append(f"Step skipped: -{SKIP_STEP_PENALTY}")

        # ── Reward formula ─────────────────────────────────────────────────────
        reward = max(0.0, base_score + step_bonus - total_penalty)

        # ── Fail trigger (individual step) ────────────────────────────────────
        fail_triggered = False
        fail_reason    = ""
        if total_penalty >= 0.8:
            fail_triggered = True
            fail_reason    = f"Single step penalty exceeded threshold ({total_penalty:.2f})"

        return StepResult(
            step            = expected_step,
            agent_response  = response,
            detected_action = detected_action,
            expected_action = expected_step.value,
            correct         = correct,
            base_score      = base_score,
            step_bonus      = step_bonus,
            penalty         = total_penalty,
            penalty_reasons = penalties,
            reward          = reward,
            fail_triggered  = fail_triggered,
            fail_reason     = fail_reason,
        )

    # ── Helper detectors ───────────────────────────────────────────────────────

    def _is_repeated_response(self, response: str) -> bool:
        if not self._response_history:
            return False
        words_new = set(response.lower().split())
        for prev in self._response_history:
            words_prev = set(prev.split())
            overlap = len(words_new & words_prev)
            if overlap >= REPEAT_WORD_OVERLAP_MIN:
                return True
        return False

    def _is_early_solution(self, response: str, step: StepName) -> bool:
        """Penalise giving resolution keywords before the resolution step."""
        if step in (StepName.EMPATHY, StepName.COLLECT_INFO):
            resolution_signals = [
                "refund", "replacement", "we will fix", "we will credit",
                "here is the solution", "the fix is", "escalate your",
            ]
            r = response.lower()
            return any(sig in r for sig in resolution_signals)
        return False

    def _is_emotion_mismatch(self, response: str) -> bool:
        """Flag cold/neutral replies when customer is angry/frustrated."""
        if self.task.customer_emotion not in ("angry", "frustrated"):
            return False
        empathy_signals = [
            "sorry", "apologize", "apology", "understand your frustration",
            "i hear you", "i completely understand", "that must be",
            "deeply sorry", "sincerely apologize",
        ]
        r = response.lower()
        return not any(sig in r for sig in empathy_signals)

    def _is_generic_response(self, response: str) -> bool:
        stripped = response.strip().lower()
        # Very short replies
        if len(stripped.split()) <= 4:
            return True
        # Generic filler phrases
        generic_phrases = [
            "ok", "okay", "done", "sure", "got it",
            "no problem", "understood", "alright",
        ]
        return stripped in generic_phrases