environment/env.py

"""
CustomerSupportEnv — OpenEnv-compliant customer support inbox environment.

Implements the standard OpenEnv interface:
  - reset(task_id, ticket_id) → Observation
  - step(action) → (Observation, Reward, done, info)
  - state() → TicketState

Three tasks:
  Task 1 (easy):   Ticket Triage
  Task 2 (medium): Guided Resolution
  Task 3 (hard):   VIP Retention & Escalation
"""
from __future__ import annotations

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

from .data import (
    TICKET_TEMPLATES, CUSTOMER_FOLLOW_UPS, KNOWLEDGE_BASE,
    get_ticket_by_id, get_random_ticket, compute_sla_deadline, time_to_sla_human,
)
from .graders import run_grader
from .models import (
    Action, ActionType, CustomerProfile, Message,
    Observation, Reward, RewardBreakdown,
    TicketCategory, TicketPriority, TicketState, TicketStatus,
)
from .rewards import compute_step_reward
from .tasks import TASK_DEFINITIONS, get_task, list_tasks


VALID_ACTIONS_BY_STATUS = {
    TicketStatus.OPEN: [
        ActionType.CLASSIFY, ActionType.RESPOND, ActionType.TAG,
        ActionType.REQUEST_INFO, ActionType.ESCALATE, ActionType.ASSIGN,
    ],
    TicketStatus.IN_PROGRESS: [
        ActionType.RESPOND, ActionType.REQUEST_INFO,
        ActionType.ESCALATE, ActionType.RESOLVE, ActionType.TAG, ActionType.ASSIGN,
    ],
    TicketStatus.WAITING_CUSTOMER: [
        ActionType.RESPOND, ActionType.ESCALATE,
        ActionType.RESOLVE, ActionType.TAG,
    ],
    TicketStatus.ESCALATED: [
        ActionType.RESPOND, ActionType.RESOLVE, ActionType.TAG,
    ],
    TicketStatus.RESOLVED: [],
    TicketStatus.CLOSED: [],
}


class CustomerSupportEnv:
    """
    OpenEnv-compliant Customer Support Inbox environment.

    Episode flow:
    1. reset(task_id, ticket_id) → agent receives Observation
    2. Agent calls step(action) repeatedly
    3. Episode ends when: ticket is Resolved/Closed, max_turns reached,
       or agent takes a RESOLVE action
    4. Final grader scores the episode (0.0–1.0)
    """

    def __init__(self, seed: Optional[int] = None):
        self._seed = seed
        if seed is not None:
            random.seed(seed)
        self._state: Optional[TicketState] = None
        self._ticket_data: Optional[Dict[str, Any]] = None
        self._task_config: Optional[Any] = None
        self._episode_history: List[Dict[str, Any]] = []
        self._cumulative_reward: float = 0.0
        self._step_count: int = 0
        self._session_id: str = str(uuid.uuid4())

    # ─── OpenEnv interface ────────────────────────────────────────────────

    def reset(
        self,
        task_id: str = "task1",
        ticket_id: Optional[str] = None,
        seed: Optional[int] = None,
    ) -> Observation:
        """
        Reset the environment for a new episode.

        Args:
            task_id: Which task to run ("task1", "task2", "task3")
            ticket_id: Specific ticket to use (None = random from task's pool)
            seed: Optional random seed for reproducibility

        Returns:
            Initial Observation for the agent
        """
        if seed is not None:
            random.seed(seed)
            self._seed = seed

        # Load task
        self._task_config = get_task(task_id)

        # Pick ticket
        if ticket_id:
            self._ticket_data = get_ticket_by_id(ticket_id)
        else:
            pool = self._task_config.ticket_pool
            chosen_id = random.choice(pool)
            self._ticket_data = get_ticket_by_id(chosen_id)

        # Build initial state
        now = datetime.now(timezone.utc)
        customer_data = self._ticket_data["customer"]
        customer = CustomerProfile(**customer_data)

        self._state = TicketState(
            ticket_id=self._ticket_data["id"],
            subject=self._ticket_data["subject"],
            body=self._ticket_data["body"],
            category=None,
            priority=None,
            status=TicketStatus.OPEN,
            tags=[],
            assigned_to=None,
            conversation=[
                Message(
                    role="customer",
                    content=self._ticket_data["body"],
                    timestamp=now.isoformat(),
                    metadata={"initial": True},
                )
            ],
            customer=customer,
            created_at=now.isoformat(),
            updated_at=now.isoformat(),
            sla_deadline=compute_sla_deadline(self._ticket_data["sla_hours"], now),
            resolution_notes=None,
            escalation_reason=None,
            turn_count=0,
            max_turns=self._task_config.max_turns,
            task_id=task_id,
            task_metadata={
                "ticket_pool": self._task_config.ticket_pool,
                "difficulty": self._task_config.difficulty,
            },
        )

        self._episode_history = []
        self._cumulative_reward = 0.0
        self._step_count = 0
        self._session_id = str(uuid.uuid4())

        return self._build_observation(last_action_result=None, last_reward=0.0)

    def step(self, action: Action) -> Tuple[Observation, Reward, bool, Dict[str, Any]]:
        """
        Execute one action in the environment.

        Args:
            action: The Action to perform

        Returns:
            (observation, reward, done, info)
        """
        if self._state is None:
            raise RuntimeError("Call reset() before step()")

        if self._state.status in (TicketStatus.RESOLVED, TicketStatus.CLOSED):
            # Episode already done
            obs = self._build_observation("Episode already ended.", 0.0)
            reward = Reward(score=0.0, cumulative_score=self._cumulative_reward,
                            done=True, breakdown=RewardBreakdown(),
                            feedback="Episode already ended.")
            return obs, reward, True, {"episode_done": True}

        # Validate action
        valid_types = VALID_ACTIONS_BY_STATUS.get(self._state.status, [])
        if action.action_type not in valid_types:
            # Invalid action — small penalty, return same observation
            obs = self._build_observation(
                f"Invalid action '{action.action_type}' for status '{self._state.status}'.",
                -0.05,
            )
            reward = Reward(
                score=0.0,
                cumulative_score=self._cumulative_reward,
                done=False,
                breakdown=RewardBreakdown(step_penalty=-0.05),
                feedback=f"Invalid action: '{action.action_type}' not allowed when status='{self._state.status}'.",
            )
            return obs, reward, False, {"error": "invalid_action"}

        prev_state = deepcopy(self._state)
        now = datetime.now(timezone.utc)

        # ── Apply action to state ──────────────────────────────────────────
        action_result = self._apply_action(action, now)

        # Increment counters
        self._state.turn_count += 1
        self._state.updated_at = now.isoformat()
        self._step_count += 1

        # Record in history
        history_entry: Dict[str, Any] = {
            "step": self._step_count,
            "timestamp": now.isoformat(),
            "action_type": action.action_type.value,
            "action_detail": action.model_dump(exclude_none=True),
            "state_status": self._state.status.value,
        }
        self._episode_history.append(history_entry)

        # ── Check episode termination ─────────────────────────────────────
        done = self._check_done()

        # ── Compute reward ────────────────────────────────────────────────
        final_grader_score = None
        if done:
            final_grader_score, breakdown, grader_feedback = run_grader(
                task_id=self._state.task_id,
                ticket_data=self._ticket_data,
                ticket_state=self._state,
                episode_history=self._episode_history,
                task_config={
                    "grader_config": self._task_config.grader_config,
                    "task_id": self._state.task_id,
                },
            )
            action_result += f" | Grader: {grader_feedback}"

        reward = compute_step_reward(
            action=action,
            prev_state=prev_state,
            new_state=self._state,
            ticket_data=self._ticket_data,
            episode_history=self._episode_history,
            done=done,
            final_score=final_grader_score,
        )

        self._cumulative_reward += reward.score
        # Update cumulative in reward object
        reward = reward.model_copy(update={
            "cumulative_score": round(min(self._cumulative_reward / max(self._step_count, 1), 1.0), 3),
            "task_complete": done and (final_grader_score or 0.0) >= self._task_config.min_score_to_pass,
        })

        obs = self._build_observation(action_result, reward.score)

        info: Dict[str, Any] = {
            "episode_id": self._session_id,
            "step": self._step_count,
            "task_id": self._state.task_id,
            "ticket_id": self._state.ticket_id,
            "status": self._state.status.value,
            "turns_used": self._state.turn_count,
            "turns_max": self._state.max_turns,
        }
        if done and final_grader_score is not None:
            info["final_grader_score"] = final_grader_score
            info["passed"] = final_grader_score >= self._task_config.min_score_to_pass

        return obs, reward, done, info

    def state(self) -> TicketState:
        """Return the current full ticket state."""
        if self._state is None:
            raise RuntimeError("Call reset() before state()")
        return deepcopy(self._state)

    # ─── Action application ───────────────────────────────────────────────

    def _apply_action(self, action: Action, now: datetime) -> str:
        """Mutate self._state based on the action. Returns human-readable result."""
        state = self._state

        if action.action_type == ActionType.CLASSIFY:
            if action.category:
                state.category = action.category
            if action.priority:
                state.priority = action.priority
            if state.status == TicketStatus.OPEN:
                state.status = TicketStatus.IN_PROGRESS
            parts = []
            if action.category:
                parts.append(f"category={action.category.value}")
            if action.priority:
                parts.append(f"priority={action.priority.value}")
            return f"Ticket classified: {', '.join(parts)}."

        elif action.action_type == ActionType.RESPOND:
            text = (action.response_text or "").strip()
            if not text:
                return "RESPOND action requires response_text."
            state.conversation.append(Message(
                role="agent",
                content=text,
                timestamp=now.isoformat(),
                metadata={"action_type": "respond"},
            ))
            if state.status == TicketStatus.OPEN:
                state.status = TicketStatus.IN_PROGRESS
            # Simulate customer follow-up if available
            self._maybe_add_customer_followup(now)
            return f"Response sent to customer ({len(text)} chars)."

        elif action.action_type == ActionType.REQUEST_INFO:
            text = (action.response_text or "").strip()
            if not text:
                return "REQUEST_INFO action requires response_text."
            state.conversation.append(Message(
                role="agent",
                content=text,
                timestamp=now.isoformat(),
                metadata={"action_type": "request_info"},
            ))
            state.status = TicketStatus.WAITING_CUSTOMER
            # Simulate customer response
            self._maybe_add_customer_followup(now, force=True)
            return f"Information requested from customer. Status → WAITING_CUSTOMER."

        elif action.action_type == ActionType.ESCALATE:
            reason = action.escalation_reason or "Escalated by agent."
            team = action.escalation_team or "tier2"
            state.escalation_reason = reason
            state.assigned_to = team
            state.status = TicketStatus.ESCALATED
            state.conversation.append(Message(
                role="agent",
                content=f"[Escalated to {team}]: {reason}",
                timestamp=now.isoformat(),
                metadata={"action_type": "escalate", "team": team},
            ))
            return f"Ticket escalated to {team}: {reason[:60]}..."

        elif action.action_type == ActionType.RESOLVE:
            notes = (action.resolution_notes or "").strip()
            if not notes:
                return "RESOLVE action requires resolution_notes."
            state.resolution_notes = notes
            state.status = TicketStatus.RESOLVED
            state.conversation.append(Message(
                role="agent",
                content=f"[Resolved]: {notes}",
                timestamp=now.isoformat(),
                metadata={
                    "action_type": "resolve",
                    "resolution_category": action.resolution_category or "unknown",
                },
            ))
            return f"Ticket resolved. Notes: {notes[:80]}..."

        elif action.action_type == ActionType.TAG:
            new_tags = action.tags or []
            # Deduplicate
            for t in new_tags:
                if t not in state.tags:
                    state.tags.append(t.lower().replace(" ", "-"))
            return f"Tags added: {new_tags}. Total tags: {state.tags}."

        elif action.action_type == ActionType.ASSIGN:
            state.assigned_to = action.assigned_to
            if state.status == TicketStatus.OPEN:
                state.status = TicketStatus.IN_PROGRESS
            return f"Ticket assigned to: {action.assigned_to}."

        return f"Action {action.action_type} processed."

    def _maybe_add_customer_followup(self, now: datetime, force: bool = False):
        """Optionally inject a simulated customer follow-up message."""
        tid = self._state.ticket_id
        followups = CUSTOMER_FOLLOW_UPS.get(tid, [])
        if not followups:
            return

        # Count existing customer follow-ups (beyond initial message)
        existing_followups = sum(
            1 for m in self._state.conversation
            if m.role == "customer" and not m.metadata.get("initial")
        )

        if existing_followups < len(followups):
            if force or (random.random() < 0.6):
                msg_text = followups[existing_followups]
                self._state.conversation.append(Message(
                    role="customer",
                    content=msg_text,
                    timestamp=now.isoformat(),
                    metadata={"followup_index": existing_followups},
                ))

    def _check_done(self) -> bool:
        """Check if the episode should end."""
        if self._state is None:
            return True
        if self._state.status in (TicketStatus.RESOLVED, TicketStatus.CLOSED):
            return True
        if self._state.turn_count >= self._state.max_turns:
            return True
        return False

    # ─── Observation builder ──────────────────────────────────────────────

    def _build_observation(
        self,
        last_action_result: Optional[str],
        last_reward: float,
    ) -> Observation:
        """Construct the Observation object from current state."""
        s = self._state
        task = self._task_config

        available = [
            a.value for a in VALID_ACTIONS_BY_STATUS.get(s.status, [])
        ]

        time_to_sla = time_to_sla_human(s.sla_deadline)

        return Observation(
            ticket_id=s.ticket_id,
            subject=s.subject,
            body=s.body,
            status=s.status,
            category=s.category,
            priority=s.priority,
            tags=s.tags,
            customer=s.customer,
            conversation=s.conversation,
            turn_count=s.turn_count,
            turns_remaining=max(0, s.max_turns - s.turn_count),
            time_to_sla=time_to_sla,
            task_id=s.task_id,
            task_description=task.description,
            task_objectives=task.objectives,
            available_actions=available,
            last_action_result=last_action_result,
            last_reward=last_reward,
            episode_done=self._check_done(),
            info={
                "difficulty": task.difficulty,
                "session_id": self._session_id,
                "step": self._step_count,
                "customer_tier": s.customer.account_tier,
                "knowledge_base": list(KNOWLEDGE_BASE.keys()),
            },
        )

    # ─── Helper methods ───────────────────────────────────────────────────

    def list_tasks(self) -> List[Dict[str, Any]]:
        """Return all available task definitions."""
        return list_tasks()

    def get_knowledge_article(self, key: str) -> str:
        """Retrieve a knowledge base article."""
        return KNOWLEDGE_BASE.get(key, "Article not found.")

    def get_episode_summary(self) -> Dict[str, Any]:
        """Return a summary of the current/last episode."""
        if self._state is None:
            return {}
        return {
            "session_id": self._session_id,
            "task_id": self._state.task_id,
            "ticket_id": self._state.ticket_id,
            "status": self._state.status.value,
            "turns_used": self._state.turn_count,
            "cumulative_reward": round(self._cumulative_reward, 3),
            "history_length": len(self._episode_history),
            "conversation_length": len(self._state.conversation),
        }