app/environment.py

"""
Session manager for the SRE Incident Response environment.
Manages per-session state, step processing, and reward accumulation.
"""

import uuid
import copy
from typing import Dict, Optional, Tuple
from datetime import datetime, timezone

from app.models import (
    Observation, Action, Reward, StepResponse,
    StateResponse, GraderResponse,
)
from app.tasks.base import AVAILABLE_ACTIONS
from app.tasks import TASK_REGISTRY


class Session:
    """In-memory session state for one episode."""

    def __init__(self, task_id: str, seed: int = 42):
        self.session_id = str(uuid.uuid4())
        self.task_id = task_id
        self.step = 0
        self.done = False
        self.total_reward = 0.0
        self.action_history: list = []
        self.created_at = datetime.now(timezone.utc).isoformat()

        task = TASK_REGISTRY[task_id]
        self.world_state = task.initial_state(seed=seed)

    def to_state_response(self, grader_score: Optional[float] = None) -> StateResponse:
        return StateResponse(
            session_id=self.session_id,
            task_id=self.task_id,
            step=self.step,
            done=self.done,
            total_reward=round(self.total_reward, 4),
            world_state=_serialize_state(self.world_state),
            action_history=self.action_history,
            grader_score=grader_score,
        )


def _serialize_state(state: dict) -> dict:
    """Make world_state JSON-serializable (convert sets, etc.)."""
    result = {}
    for k, v in state.items():
        if k.startswith("_"):
            continue  # hide internal fields
        if isinstance(v, set):
            result[k] = list(v)
        elif isinstance(v, dict):
            result[k] = _serialize_state(v)
        else:
            result[k] = v
    return result


class EnvironmentManager:
    """Manages all active sessions."""

    def __init__(self):
        self._sessions: Dict[str, Session] = {}

    @staticmethod
    def _clamp_score_strict(score: float, eps: float = 0.01) -> float:
        """
        Hackathon validator requirement: task scores must be strictly within (0, 1).
        We clamp away from exact endpoints to avoid returning 0.0 or 1.0.
        """
        try:
            s = float(score)
        except Exception:
            s = 0.0
        if s <= 0.0:
            return eps
        if s >= 1.0:
            return 1.0 - eps
        return s

    def reset(self, task_id: str, seed: int = 42) -> Tuple[Observation, str]:
        """Start a new episode. Returns (initial_observation, session_id)."""
        if task_id not in TASK_REGISTRY:
            raise ValueError(f"Unknown task_id '{task_id}'. Valid: {list(TASK_REGISTRY.keys())}")

        session = Session(task_id=task_id, seed=seed)
        self._sessions[session.session_id] = session

        task = TASK_REGISTRY[task_id]
        obs = task.get_observation(session.world_state, session.session_id, step=0)
        obs.message = (
            f"New episode started. Task: {task.name} (difficulty: {task.difficulty}). "
            f"Max steps: {task.max_steps}. Investigate the incident and resolve it."
        )
        return obs, session.session_id

    def step(self, session_id: str, action: Action) -> StepResponse:
        """Process one action and return (observation, reward, done, info)."""
        if session_id not in self._sessions:
            raise KeyError(f"Session '{session_id}' not found. Call /reset first.")

        session = self._sessions[session_id]

        if session.done:
            task = TASK_REGISTRY[session.task_id]
            obs = task.get_observation(session.world_state, session_id, session.step)
            obs.message = "Episode already complete. Call /reset to start a new episode."
            return StepResponse(
                observation=obs,
                reward=Reward(value=0.0, cumulative=session.total_reward, breakdown={},
                              message="Episode already complete."),
                done=True,
                info={"episode_complete": True},
            )

        task = TASK_REGISTRY[session.task_id]
        session.step += 1

        # Check max steps
        if session.step > task.max_steps:
            session.done = True
            obs = task.get_observation(session.world_state, session_id, session.step)
            obs.message = f"Episode ended: max steps ({task.max_steps}) reached without resolution."
            timeout_reward = -0.10
            session.total_reward += timeout_reward
            return StepResponse(
                observation=obs,
                reward=Reward(
                    value=timeout_reward,
                    cumulative=round(session.total_reward, 4),
                    breakdown={"timeout_penalty": timeout_reward},
                    message="Max steps reached.",
                ),
                done=True,
                info={"timeout": True, "steps": session.step},
            )

        # Process action
        new_state, step_reward, done, message = task.process_action(
            action.action_type,
            action.parameters,
            session.world_state,
        )
        session.world_state = new_state
        session.done = done
        session.total_reward += step_reward

        # Record history
        session.action_history.append({
            "step": session.step,
            "action_type": action.action_type,
            "parameters": action.parameters,
            "reward": round(step_reward, 4),
            "message_preview": message[:120] if message else "",
        })

        # Build observation
        obs = task.get_observation(session.world_state, session_id, session.step)
        obs.message = message

        # Grade for info
        score, grade_breakdown = task.grade(session.world_state, session.action_history)
        score = self._clamp_score_strict(score)

        reward_obj = Reward(
            value=round(step_reward, 4),
            cumulative=round(session.total_reward, 4),
            breakdown=_build_reward_breakdown(action.action_type, step_reward),
            message=f"Step reward: {step_reward:+.3f} | Cumulative: {session.total_reward:+.3f}",
        )

        return StepResponse(
            observation=obs,
            reward=reward_obj,
            done=done,
            info={
                "step": session.step,
                "max_steps": task.max_steps,
                "grader_score": score,
                "episode_complete": done,
            },
        )

    def get_state(self, session_id: str) -> StateResponse:
        if session_id not in self._sessions:
            raise KeyError(f"Session '{session_id}' not found.")
        session = self._sessions[session_id]
        task = TASK_REGISTRY[session.task_id]
        score, _ = task.grade(session.world_state, session.action_history)
        score = self._clamp_score_strict(score)
        return session.to_state_response(grader_score=score)

    def grade(self, session_id: str) -> GraderResponse:
        if session_id not in self._sessions:
            raise KeyError(f"Session '{session_id}' not found.")
        session = self._sessions[session_id]
        task = TASK_REGISTRY[session.task_id]
        score, breakdown = task.grade(session.world_state, session.action_history)
        score = self._clamp_score_strict(score)

        passing = score >= task.passing_score
        return GraderResponse(
            session_id=session_id,
            task_id=session.task_id,
            score=score,
            breakdown=breakdown,
            episode_complete=session.done,
            steps_taken=session.step,
            message=(
                f"Score: {score:.4f} | "
                f"{'PASS' if passing else 'FAIL'} "
                f"(threshold: {task.passing_score}) | "
                f"Steps: {session.step}/{task.max_steps}"
            ),
        )

    def cleanup_session(self, session_id: str) -> bool:
        if session_id in self._sessions:
            del self._sessions[session_id]
            return True
        return False

    def active_sessions(self) -> int:
        return len(self._sessions)


def _build_reward_breakdown(action_type: str, value: float) -> Dict[str, float]:
    if value > 0:
        return {f"{action_type}_reward": round(value, 4)}
    elif value < 0:
        return {f"{action_type}_penalty": round(value, 4)}
    return {}