# DebugOps Environment — core RL environment for AI incident response.
# Follows the OpenEnv step()/reset()/state() interface.
from __future__ import annotations
import copy
from dataclasses import dataclass, field, asdict
from typing import Dict, List, Tuple, Any

from env.incident_generator import generate_incident
from env.dynamics import apply_action
from env.reward import compute_reward

# Typed observation / state model
@dataclass
class Observation:
    services: Dict[str, str]          # service_name -> "healthy" | "degraded"
    logs: List[str]                   # ordered log lines (may be noisy)
    metrics: Dict[str, float]         # latency, error_rate, cpu
    time_step: int                    # steps elapsed in the current episode
    fix_progress: int                 # number of correct steps completed so far
    metric_trend: str                 # "improving" | "degrading" | "stable"

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


# Core environment
class DebugEnv:
    """
    Single-episode production incident environment.

    Observation space:
        services    : Dict[str, str]   — per-service health status
        logs        : List[str]        — system logs (may contain noise)
        metrics     : Dict[str,float]  — latency (ms), error_rate (0-1), cpu (%)
        time_step   : int

    Action space (discrete, 5 actions):
        restart_api | restart_db | restart_cache | scale_up | noop

    Episode terminates when:
        - state_data["resolved"] is True  (success)
        - t >= max_steps                  (timeout / failure)
    """

    VALID_ACTIONS = ["restart_api", "restart_db", "restart_cache", "scale_up", "noop"]

    def __init__(self, max_steps: int = 20):
        self.max_steps = max_steps
        self.t: int = 0
        self.done: bool = False
        self.success: bool = False
        self.state_data: Dict[str, Any] = {}
        self._prev_latency: float = 0.0

    #---
    def reset(self) -> Dict[str, Any]:
        """Return fresh observation; episode counter reset."""
        self.t = 0
        self.done = False
        self.success = False
        self.state_data = generate_incident()
        self._prev_latency = self.state_data["metrics"]["latency"]
        return self._obs()

    #---
    def state(self) -> Dict[str, Any]:
        """Return current observation (idempotent)."""
        return self._obs()

    
    def step(self, action: str) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]:
        """
        Apply action and advance one time-step.

        Returns
        -------
        observation : Dict
        reward      : float
        done        : bool
        info        : Dict   — latency, error_rate, progress, resolved, success
        """
        if self.done:
            raise RuntimeError("Episode is finished. Call reset() before stepping.")

        if action not in self.VALID_ACTIONS:
            raise ValueError(f"Invalid action '{action}'. Must be one of {self.VALID_ACTIONS}")

        prev_state = copy.deepcopy(self.state_data)
        prev_latency = self.state_data["metrics"]["latency"]
        self.state_data = apply_action(self.state_data, action)
        reward, info = compute_reward(prev_state, self.state_data, action, self.t)
        self._prev_latency = prev_latency

        self.t += 1

        if self.state_data["resolved"]:
            self.done = True
            self.success = True
        elif self.t >= self.max_steps:
            self.done = True
            self.success = False

        info["success"] = self.success
        info["time_step"] = self.t

        return self._obs(), reward, self.done, info

    
    def _obs(self) -> Dict[str, Any]:
        curr_latency = self.state_data["metrics"]["latency"]
        delta = curr_latency - self._prev_latency
        if delta < -10:
            trend = "improving"
        elif delta > 10:
            trend = "degrading"
        else:
            trend = "stable"

        return Observation(
            services=dict(self.state_data["services"]),
            logs=list(self.state_data["logs"]),
            metrics={k: round(v, 2) for k, v in self.state_data["metrics"].items()},
            time_step=self.t,
            fix_progress=self.state_data["fix_progress"],
            metric_trend=trend,
        ).to_dict()