server/environment.py

"""
server/environment.py — Core OpenEnv environment for Cloud Incident Response.

Difficulty comes from SCENARIO DESIGN, not mechanics:
  EASY:   3 services, clear metrics, obvious severity
  MEDIUM: 8 services, root cause NOT in alert, must follow log breadcrumbs
  HARD:   8 services + 5-7 remediation steps + quality summary + penalties
"""

from __future__ import annotations

import os
import sys
import threading
import uuid

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from graders import _svc_match, grade
from server.models import Action, ActionParameters, EpisodeState, Observation, Reward
from tasks import get_scenario, get_task

_DIAGNOSTIC = frozenset({
    "query_logs", "check_metrics", "check_dependencies",
    "check_recent_deploys", "check_service_status",
})

_REMEDIATION = frozenset({
    "restart_service", "rollback_deploy", "scale_service",
    "disable_feature_flag", "clear_cache", "execute_runbook_step",
})

_SUBMIT = frozenset({
    "submit_severity", "submit_root_cause", "submit_resolution",
})

_TASK_SUBMIT = {
    "alert_classification": "submit_severity",
    "root_cause_analysis": "submit_root_cause",
    "remediation_planning": "submit_resolution",
}

_REWARD_TABLE = {
    "easy": {
        "query_new_svc":     +0.04,  "query_new_action":  +0.02,
        "query_repeat":      -0.03,  "query_unknown_svc": -0.06,
        "query_no_service":  -0.04,  "rem_good":          +0.00,
        "rem_wrong":         -0.08,  "rem_no_target":     -0.05,
        "submit_correct":    +0.02,  "submit_wrong":      -0.08,
        "past_half":         -0.04,  "timeout":           -0.15,
        "bad_action":        -0.05,  "exact_repeat":      -0.04,
    },
    "medium": {
        "query_new_svc":     +0.04,  "query_new_action":  +0.02,
        "query_repeat":      -0.04,  "query_unknown_svc": -0.06,
        "query_no_service":  -0.04,  "rem_good":          +0.06,
        "rem_wrong":         -0.10,  "rem_no_target":     -0.06,
        "submit_correct":    +0.02,  "submit_wrong":      -0.10,
        "past_half":         -0.02,  "timeout":           -0.15,
        "bad_action":        -0.05,  "exact_repeat":      -0.05,
    },
    "hard": {
        "query_new_svc":     +0.03,  "query_new_action":  +0.01,
        "query_repeat":      -0.03,  "query_unknown_svc": -0.05,
        "query_no_service":  -0.03,  "rem_good":          +0.06,
        "rem_wrong":         -0.15,  "rem_no_target":     -0.05,
        "submit_correct":    +0.02,  "submit_wrong":      -0.12,
        "past_half":         -0.02,  "timeout":           -0.20,
        "bad_action":        -0.05,  "exact_repeat":      -0.04,
    },
}

_TASK_DIFFICULTY = {
    "alert_classification": "easy",
    "root_cause_analysis": "medium",
    "remediation_planning": "hard",
}


class IncidentEnvironment:
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._s: dict = {}
        self._scenario: dict = {}
        self._task_def: dict = {}
        self._ready = False

    def reset(self, task_id: str = "alert_classification",
              scenario_index: int = 0) -> Observation:
        with self._lock:
            task_def = get_task(task_id)
            scenario = get_scenario(task_id, scenario_index)
            self._task_def = task_def
            self._scenario = scenario
            self._s = {
                "episode_id": str(uuid.uuid4()),
                "task_id": task_id,
                "scenario_id": scenario["scenario_id"],
                "step_count": 0,
                "max_steps": task_def["max_steps"],
                "action_history": [],
                "queried_data": {},
                "queried_keys": set(),
                "services_queried": set(),
                "exact_hashes": set(),
                "submitted": False,
                "resolved": False,
                "done": False,
                "cumulative_reward": 0.0,
                "feedback": f"Episode started. {scenario['description']}",
                "last_action_error": None,
            }
            self._ready = True
            return self._build_obs()

    def step(self, action: Action) -> tuple[Observation, Reward, bool, dict]:
        with self._lock:
            if not self._ready:
                raise RuntimeError("Call reset() before step().")
            s = self._s
            s["last_action_error"] = None

            if s["done"]:
                return (self._build_obs(),
                        Reward(score=0.0, reason="episode already done",
                               cumulative=s["cumulative_reward"]),
                        True, {})

            s["step_count"] += 1
            step_num = s["step_count"]
            at = action.action_type
            params = action.parameters
            task_id = s["task_id"]
            diff = _TASK_DIFFICULTY.get(task_id, "medium")
            rt = _REWARD_TABLE[diff]

            s["action_history"].append({
                "action_type": at,
                "parameters": params.model_dump(exclude_none=True),
                "step": step_num,
            })

            r = 0.0
            fb: list[str] = []

            h = f"{at}|{params.model_dump_json(exclude_none=True)}"
            if h in s["exact_hashes"]:
                r += rt["exact_repeat"]
                fb.append(f"exact repeat ({rt['exact_repeat']:+.2f})")
            s["exact_hashes"].add(h)

            half = max(1, s["max_steps"] // 2)
            if step_num > half and at not in _SUBMIT:
                r += rt["past_half"]
                fb.append(f"past halfway ({rt['past_half']:+.3f})")

            if at in _DIAGNOSTIC:
                r, fb = self._handle_diagnostic(at, params, r, fb, rt)
            elif at in _REMEDIATION:
                r, fb = self._handle_remediation(at, params, r, fb, rt, task_id)
            elif at in _SUBMIT:
                r, fb, terminal = self._handle_submit(at, params, r, fb, rt, task_id)
                if terminal:
                    s["done"] = True
            else:
                r += rt["bad_action"]
                fb.append(f"unknown action '{at}' ({rt['bad_action']:+.2f})")
                s["last_action_error"] = f"Unknown action type: {at}"

            if step_num >= s["max_steps"] and not s["done"]:
                r += rt["timeout"]
                fb.append(f"timeout ({rt['timeout']:+.2f})")
                s["done"] = True

            if s["done"]:
                result = grade(s["task_id"], s, self._scenario)
                grader_score = result["total"]
                s["cumulative_reward"] = round(
                    s["cumulative_reward"] + r + grader_score, 4)
                fb.append(f"grader={grader_score:.3f} ({result['feedback']})")
            else:
                s["cumulative_reward"] = round(s["cumulative_reward"] + r, 4)

            s["feedback"] = " | ".join(fb) if fb else "ok"
            return (self._build_obs(),
                    Reward(score=round(r, 4), reason=s["feedback"],
                           cumulative=s["cumulative_reward"]),
                    s["done"],
                    {"step": step_num, "feedback": s["feedback"]})

    def state(self) -> EpisodeState:
        with self._lock:
            if not self._ready:
                raise RuntimeError("No active episode — call reset() first.")
            s = self._s
            return EpisodeState(
                episode_id=s["episode_id"], task_id=s["task_id"],
                scenario_id=s["scenario_id"], step_count=s["step_count"],
                max_steps=s["max_steps"],
                action_history=list(s["action_history"]),
                queried_data=dict(s["queried_data"]),
                submitted=s["submitted"], resolved=s["resolved"],
                done=s["done"], cumulative_reward=s["cumulative_reward"],
                feedback=s["feedback"])

    def _handle_diagnostic(self, at, params, r, fb, rt):
        s = self._s
        svc = (params.service or "").lower().strip()
        known = {v.lower() for v in self._scenario.get("known_services", set())}
        tool = self._scenario.get("tool_responses", {}).get(at, {})
        key = (at, svc)

        if not svc:
            r += rt["query_no_service"]
            fb.append(f"{at}: no service ({rt['query_no_service']:+.2f})")
            s["last_action_error"] = f"{at} requires a service parameter"
            return r, fb

        if svc not in known:
            r += rt["query_unknown_svc"]
            fb.append(f"unknown service '{svc}' ({rt['query_unknown_svc']:+.2f})")
            s["last_action_error"] = f"Unknown service: {svc}"
            return r, fb

        if key in s["queried_keys"]:
            r += rt["query_repeat"]
            fb.append(f"repeat [{at}][{svc}] ({rt['query_repeat']:+.2f})")
        elif svc in s["services_queried"]:
            r += rt["query_new_action"]
            fb.append(f"new action on {svc} ({rt['query_new_action']:+.2f})")
            s["queried_keys"].add(key)
        else:
            r += rt["query_new_svc"]
            fb.append(f"new service {svc} ({rt['query_new_svc']:+.2f})")
            s["queried_keys"].add(key)
            s["services_queried"].add(svc)

        result = tool.get(svc, f"No data available for '{svc}'.")
        s["queried_data"].setdefault(at, {})[svc] = result
        return r, fb

    def _handle_remediation(self, at, params, r, fb, rt, task_id):
        s = self._s
        if task_id == "alert_classification":
            r += rt["rem_wrong"]
            fb.append(f"remediation in easy task ({rt['rem_wrong']:+.2f})")
            s["last_action_error"] = "Remediation not available in alert_classification"
            return r, fb

        svc = (params.service or "").lower().strip()
        flag = (params.flag or "").lower().strip()
        runbook = (params.runbook_action or "").lower().strip()
        target = (params.target or params.target_version or "").lower().strip()

        if not (svc or flag or runbook or target):
            r += rt["rem_no_target"]
            fb.append(f"{at}: no target ({rt['rem_no_target']:+.2f})")
            s["last_action_error"] = f"{at} requires a target"
            return r, fb

        keys = {at}
        if svc:     keys.add(f"{at}:{svc}")
        if flag:    keys.add(f"{at}:{flag}")
        if runbook: keys.add(f"execute_runbook_step:{runbook}")
        if target:  keys.add(f"execute_runbook_step:{target}")

        wrong_map = self._scenario.get("wrong_actions", {})
        rem_data = self._scenario.get("remediation_data", {})

        is_wrong = any(k in wrong_map for k in keys)
        if not is_wrong and svc:
            for wk in wrong_map:
                if ":" in wk:
                    w_at, w_svc = wk.split(":", 1)
                    if w_at == at and _svc_match(svc, w_svc):
                        is_wrong = True
                        break

        if is_wrong:
            r += rt["rem_wrong"]
            reason = next((wrong_map[k] for k in keys if k in wrong_map), "wrong")
            fb.append(f"wrong: {at} — {str(reason)[:60]} ({rt['rem_wrong']:+.2f})")
        else:
            r += rt["rem_good"]
            tgt = svc or flag or runbook or target
            fb.append(f"executed {at}:{tgt} ({rt['rem_good']:+.2f})")
            at_data = rem_data.get(at, {})
            result = (at_data.get(svc) or at_data.get(flag) or at_data.get(runbook)
                      or at_data.get(target) or "action executed successfully")
            s["queried_data"].setdefault(at, {})[tgt] = result
        return r, fb

    def _handle_submit(self, at, params, r, fb, rt, task_id):
        s = self._s
        correct = _TASK_SUBMIT.get(task_id, "")
        if at != correct:
            r += rt["submit_wrong"]
            fb.append(f"wrong submit '{at}' (need '{correct}') ({rt['submit_wrong']:+.2f})")
            s["last_action_error"] = f"Wrong submission type: use {correct}"
            return r, fb, False

        s["submitted"] = True
        r += rt["submit_correct"]
        fb.append(f"submitted ({rt['submit_correct']:+.2f})")

        if at == "submit_severity":
            fb.append(f"severity={(params.severity or '').upper().strip()}")
        elif at == "submit_root_cause":
            fb.append(f"svc={params.service or ''}, mode={params.failure_mode or ''}")
        elif at == "submit_resolution":
            summary = params.summary or ""
            inv = sum(1 for a in s["action_history"]
                      if a.get("action_type") in _DIAGNOSTIC | _REMEDIATION)
            if summary.strip() and inv >= 1:
                s["resolved"] = True
                fb.append("resolved")
            else:
                fb.append("insufficient investigation")
        return r, fb, True

    def _build_obs(self):
        s = self._s
        sc = self._scenario
        td = self._task_def
        return Observation(
            episode_id=s["episode_id"], task_id=s["task_id"],
            scenario_id=s["scenario_id"], step_count=s["step_count"],
            max_steps=s["max_steps"],
            incident_summary=sc.get("incident_summary", sc.get("description", "")),
            alert=sc.get("alert", {}),
            available_actions=td.get("available_actions", []),
            queried_data=dict(s["queried_data"]),
            cumulative_reward=s["cumulative_reward"],
            done=s["done"], feedback=s["feedback"],
            known_services=sorted(sc.get("known_services", set())),
            last_action_error=s.get("last_action_error"))