# Task: Critical — SLA-breach scenario with controlled randomness and realism.
# Introduces noisy logs, metric variance, and slight ambiguity while preserving solvability.

from __future__ import annotations
from typing import Dict, Any, Tuple
import random

from env.environment import DebugEnv


class CriticalEnv(DebugEnv):
    """
    High-stakes incident:
      - Root cause: memory_leak (restart_api → restart_db)
      - Noisy + misleading logs
      - Randomized metrics per episode
      - SLA penalties + time pressure
      - Designed for score variance across runs
    """

    SLA_LATENCY_THRESHOLD = 250  # ms
    SLA_PENALTY = 80.0
    TIME_PRESSURE_START = 8
    TIME_PRESSURE_PENALTY = 30.0

    def reset(self) -> Dict[str, Any]:
        state = super().reset()

        self.state_data["root_cause"] = "memory_leak"
        self.state_data["fix_sequence"] = ["restart_api", "restart_db"]

        self.state_data["services"]["api"] = "degraded"
        self.state_data["services"]["db"] = "degraded"

        base_logs = [
            "memory usage increasing",
            "OOM warning",
            "heap allocation failure",
        ]

        noise_logs = [
            "network latency spike (transient)",
            "disk almost full: /var/log 94%",
            "temporary service restart: metrics-collector",
            "cache eviction rate elevated",
            "connection pool retry",
            "upstream request timeout",
        ]

        # Randomly inject noise
        selected_noise = random.sample(noise_logs, k=random.randint(2, 4))

        logs = base_logs + selected_noise
        random.shuffle(logs)

        self.state_data["logs"] = logs
        self.state_data["metrics"]["latency"] = random.randint(350, 500)
        self.state_data["metrics"]["error_rate"] = round(random.uniform(0.6, 0.9), 2)
        self.state_data["metrics"]["cpu"] = random.randint(85, 98)

        self._prev_latency = self.state_data["metrics"]["latency"]

        return self._obs()

    def step(self, action: str) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]:
        obs, reward, done, info = super().step(action)

        if info["latency"] > self.SLA_LATENCY_THRESHOLD:
            reward -= self.SLA_PENALTY

        if self.t > self.TIME_PRESSURE_START and not info["resolved"]:
            reward -= self.TIME_PRESSURE_PENALTY

        reward += random.uniform(-3, 3)

        return obs, reward, done, info


def create_env() -> CriticalEnv:
    return CriticalEnv(max_steps=10)