baselines.py

"""
baselines.py — Three baseline agents benchmarked across all 10 ARIA tasks.

Scoring: uses grade_episode() from graders.grader — guaranteed [0.0, 1.0].
Seeds: each task runs 20 different seeds for variance measurement.

Agents:
  1. RandomAgent         — random ActionType + random service each step
  2. RuleBasedAgent      — heuristic: find worst service, read -> diagnose -> fix
  3. InformationFirstAgent — gather all evidence -> pick best fix

Usage:
  python baselines.py
"""
from __future__ import annotations
import json
import random
import statistics
from datetime import datetime, timezone
from typing import List, Dict, Optional, Tuple

from env import DevOpsIncidentEnv
from models import Action, ActionType, Observation, StepResult
from graders.grader import grade_episode

# --- Constants ----------------------------------------------------------------

ALL_TASKS = [
    "easy", "medium", "hard", "bonus",
    "security", "database", "failover",
    "dns", "waf", "thundering_herd",
]

NUM_SEEDS = 20
RUNBOOKS = [
    "high_cpu.md", "memory_leak.md", "db_connection.md",
    "deployment_rollback.md", "cascade_failure.md", "data_corruption.md",
    "dns_failure.md", "waf_cpu.md", "thundering_herd.md",
]

# Keywords that help InformationFirstAgent pick the right runbook
RUNBOOK_KEYWORDS = {
    "dns_failure.md":        ["dns", "nxdomain", "nameserver", "zone", "puppet"],
    "waf_cpu.md":            ["waf", "regex", "backtrack", "cpu", "rule", "firewall"],
    "thundering_herd.md":    ["cache", "thundering", "herd", "vitess", "redis"],
    "memory_leak.md":        ["memory", "oom", "heap", "leak", "outofmemory"],
    "db_connection.md":      ["database", "db", "connection", "pool", "query", "postgres"],
    "high_cpu.md":           ["cpu", "spike", "thread", "exhaustion"],
    "deployment_rollback.md":["deploy", "rollback", "version", "manifest"],
    "cascade_failure.md":    ["cascade", "circuit breaker", "upstream", "downstream"],
    "data_corruption.md":    ["corrupt", "corruption", "mismatch", "silent", "data"],
}


# --- Episode runner -----------------------------------------------------------

def run_episode(task_id: str, seed: int, agent_fn) -> float:
    """
    Run one full episode and return a grade_episode() score in [0.0, 1.0].
    The seed is passed to both env construction AND env.reset() for full
    seed-based variation.
    """
    env = DevOpsIncidentEnv(task_id=task_id, seed=seed)
    obs = env.reset(seed=seed)
    done = False
    while not done:
        action = agent_fn(obs)
        result: StepResult = env.step(action)
        done = result.done
        obs = result.observation

    s = env.state()
    score = grade_episode(
        task_id=task_id,
        action_history=s.action_history,
        ground_truth_root_cause=s.ground_truth_root_cause,
        ground_truth_fix=s.ground_truth_fix,
        incident_resolved=s.incident_resolved,
        total_reward=s.total_reward,
    )
    # Guaranteed [0.0, 1.0] by grader contract, but assert for safety
    assert 0.0 <= score <= 1.0, f"grade_episode() returned out-of-range score {score} for {task_id} seed={seed}"
    return round(score, 4)


# --- Helpers ------------------------------------------------------------------

def _worst_service(obs: Observation) -> Optional[str]:
    """Return name of the most-degraded service (highest error_rate)."""
    down = [s for s in obs.services if s.status == "down"]
    degraded = [s for s in obs.services if s.status == "degraded"]
    pool = down or degraded
    if not pool:
        return obs.services[0].name if obs.services else None
    return max(pool, key=lambda s: s.error_rate).name


def _unhealthy_services(obs: Observation) -> List[str]:
    return [s.name for s in obs.services if s.status in ("down", "degraded")]


def _evidence_text(obs: Observation) -> str:
    return " ".join(e.raw.lower() for e in obs.evidence_log)


def _pick_runbook(evidence: str) -> Optional[str]:
    best_rb, best_hits = None, 0
    for rb, keywords in RUNBOOK_KEYWORDS.items():
        hits = sum(1 for kw in keywords if kw in evidence)
        if hits > best_hits:
            best_hits, best_rb = hits, rb
    return best_rb


# --- Agent 1: RandomAgent -----------------------------------------------------

class RandomAgent:
    """
    Picks a random ActionType and (where applicable) a random service.
    Pure lower-bound baseline. Uses per-seed RNG so each episode varies.
    """
    SERVICE_ACTIONS = {
        ActionType.READ_LOGS, ActionType.READ_METRICS, ActionType.SEARCH_LOGS,
        ActionType.RESTART_SERVICE, ActionType.ROLLBACK, ActionType.SCALE_UP,
        ActionType.ACKNOWLEDGE,
    }

    def __init__(self, seed: int = 0):
        # Different offset per seed so episodes differ
        self.rng = random.Random(seed * 1000 + 7777)

    def __call__(self, obs: Observation) -> Action:
        at = self.rng.choice(list(ActionType))
        svc = None
        if at in self.SERVICE_ACTIONS and obs.services:
            svc = self.rng.choice(obs.services).name
        rb = None
        if at == ActionType.READ_RUNBOOK and RUNBOOKS:
            rb = self.rng.choice(RUNBOOKS)
        rc = "unknown_failure" if at == ActionType.DIAGNOSE else None
        return Action(action_type=at, service=svc, runbook=rb, root_cause=rc)


# --- Agent 2: RuleBasedAgent --------------------------------------------------

class RuleBasedAgent:
    """
    Heuristic three-phase pipeline:
      Phase 1 (Gather):  Read logs + metrics for worst service; pick runbook from alert
      Phase 2 (Diagnose): Diagnose as "service failure in <name>"
      Phase 3 (Fix):     Restart if down; rollback if historically dated; else alert_oncall

    Deterministic — std deviation shows cross-seed task variation (not agent variance).
    """

    def __init__(self):
        self._phase: Dict[str, str] = {}
        self._gathered: Dict[str, set] = {}
        self._target: Dict[str, Optional[str]] = {}

    def __call__(self, obs: Observation) -> Action:
        ek = obs.task_id
        if obs.step == 0 or ek not in self._phase:
            self._phase[ek] = "gather"
            self._gathered[ek] = set()
            self._target[ek] = _worst_service(obs)

        target = self._target[ek] or (obs.services[0].name if obs.services else None)
        gathered = self._gathered[ek]

        # Phase 1: Gather
        if self._phase[ek] == "gather":
            if "logs" not in gathered:
                gathered.add("logs")
                return Action(action_type=ActionType.READ_LOGS, service=target)
            if "metrics" not in gathered:
                gathered.add("metrics")
                return Action(action_type=ActionType.READ_METRICS, service=target)
            if "runbook" not in gathered:
                gathered.add("runbook")
                rb = "memory_leak.md"
                if obs.active_alerts:
                    msg = obs.active_alerts[0].message.lower()
                    if "dns" in msg or "nxdomain" in msg:
                        rb = "dns_failure.md"
                    elif "waf" in msg or "regex" in msg or "cpu 100" in msg:
                        rb = "waf_cpu.md"
                    elif "cache" in msg or "thundering" in msg or "vitess" in msg:
                        rb = "thundering_herd.md"
                    elif "database" in msg or "index" in msg or "query" in msg:
                        rb = "db_connection.md"
                    elif "deploy" in msg or "rollback" in msg:
                        rb = "deployment_rollback.md"
                    elif "ddos" in msg or "attack" in msg or "botnet" in msg:
                        rb = "cascade_failure.md"
                    elif "corrupt" in msg or "mismatch" in msg:
                        rb = "data_corruption.md"
                return Action(action_type=ActionType.READ_RUNBOOK, runbook=rb)
            self._phase[ek] = "diagnose"

        # Phase 2: Diagnose
        if self._phase[ek] == "diagnose":
            self._phase[ek] = "fix"
            return Action(
                action_type=ActionType.DIAGNOSE,
                root_cause=f"service failure in {target}",
            )

        # Phase 3: Fix
        svc_map = {s.name: s for s in obs.services}
        if target and target in svc_map:
            svc = svc_map[target]
            if svc.status == "down":
                return Action(action_type=ActionType.RESTART_SERVICE, service=target)
            if any(yr in svc.last_deployed for yr in ["2019", "2022", "2016"]):
                return Action(action_type=ActionType.ROLLBACK, service=target)
            if svc.status == "degraded":
                return Action(action_type=ActionType.RESTART_SERVICE, service=target)

        return Action(
            action_type=ActionType.ALERT_ONCALL,
            reason="service degradation requires manual ops attention",
        )


# --- Agent 3: InformationFirstAgent -------------------------------------------

class InformationFirstAgent:
    """
    Exhaustive evidence gathering before any remediation.

    Pipeline:
      1. Read logs for ALL degraded/down services (one per step)
      2. Read metrics for ALL degraded/down services (one per step)
      3. Pick best-matching runbook from accumulated evidence keywords
      4. Diagnose based on keyword signals in evidence
      5. Apply the most appropriate fix ranked by evidence signals
    """

    FIX_HEURISTICS = [
        # (evidence_keywords, action_type, extra_kwargs)
        (["puppet", "manifest", "nameserver", "zone transfer"],
         ActionType.RESTART_SERVICE, {"service_kw": "caching"}),
        (["waf", "regex", "backtrack", "sig-2019", "rule-set"],
         ActionType.ROLLBACK, {"service_kw": "waf-rule-manager"}),
        (["thundering herd", "vitess", "connection pool exhausted"],
         ActionType.RESTART_SERVICE, {"service_kw": "vitess"}),
        (["cache miss", "redis", "simultaneous"],
         ActionType.ALERT_ONCALL, {"reason": "restore redis cache nodes to resolve thundering herd"}),
        (["ddos", "botnet", "185.", "credential stuffing"],
         ActionType.BLOCK_IP_RANGE, {"ip_range": "185.0.0.0/8"}),
        (["sequential scan", "missing index", "full table scan"],
         ActionType.CREATE_INDEX, {}),
        (["oom", "outofmemory", "heap", "memory"],
         ActionType.RESTART_SERVICE, {"service_kw": None}),
        (["cascade", "connection pool", "rollback"],
         ActionType.ROLLBACK, {"service_kw": None}),
        (["partition", "region", "failover"],
         ActionType.FAILOVER, {}),
        (["corrupt", "mismatch", "price", "data"],
         ActionType.ROLLBACK, {"service_kw": None}),
    ]

    def __init__(self):
        self._state: Dict[str, dict] = {}

    def _init_state(self, obs: Observation) -> dict:
        unhealthy = _unhealthy_services(obs)
        return {
            "queue_logs":    list(unhealthy),
            "queue_metrics": list(unhealthy),
            "runbook_done":  False,
            "diagnosed":     False,
            "fixes_applied": set(),
        }

    def __call__(self, obs: Observation) -> Action:
        ek = obs.task_id
        if obs.step == 0 or ek not in self._state:
            self._state[ek] = self._init_state(obs)

        st = self._state[ek]
        svc_names = {s.name for s in obs.services}

        # Gather logs
        while st["queue_logs"]:
            svc = st["queue_logs"].pop(0)
            if svc in svc_names:
                return Action(action_type=ActionType.READ_LOGS, service=svc)

        # Gather metrics
        while st["queue_metrics"]:
            svc = st["queue_metrics"].pop(0)
            if svc in svc_names:
                return Action(action_type=ActionType.READ_METRICS, service=svc)

        # Read runbook
        if not st["runbook_done"]:
            st["runbook_done"] = True
            evidence = _evidence_text(obs)
            rb = _pick_runbook(evidence) or "high_cpu.md"
            return Action(action_type=ActionType.READ_RUNBOOK, runbook=rb)

        # Diagnose
        if not st["diagnosed"]:
            st["diagnosed"] = True
            evidence = _evidence_text(obs)
            kw_map = [
                ("puppet",           "puppet_manifest_bug"),
                ("nameserver",       "nameserver_misconfiguration"),
                ("nxdomain",         "dns_nxdomain_failure"),
                ("zone",             "dns_zone_corruption"),
                ("regex",            "waf_catastrophic_backtracking"),
                ("backtrack",        "waf_cpu_exhaustion"),
                ("thundering",       "thundering_herd"),
                ("vitess",           "vitess_connection_pool_exhaustion"),
                ("redis",            "cache_node_failure"),
                ("ddos",             "ddos_attack"),
                ("botnet",           "credential_stuffing_botnet"),
                ("index",            "missing_database_index"),
                ("sequential scan",  "full_table_scan"),
                ("oom",              "memory_oom_kill"),
                ("heap",             "heap_exhaustion"),
                ("corrupt",          "data_corruption"),
                ("cascade",          "cascade_failure"),
                ("partition",        "network_partition"),
            ]
            terms = []
            for kw, label in kw_map:
                if kw in evidence and label not in terms:
                    terms.append(label)
            rc = " + ".join(terms[:3]) if terms else "unknown_service_failure"
            return Action(action_type=ActionType.DIAGNOSE, root_cause=rc)

        # Apply fixes
        evidence = _evidence_text(obs)

        for keywords, action_type, kwargs in self.FIX_HEURISTICS:
            fix_key = "|".join(keywords)
            if fix_key in st["fixes_applied"]:
                continue
            if any(kw in evidence for kw in keywords):
                st["fixes_applied"].add(fix_key)
                if "service_kw" in kwargs:
                    kw = kwargs["service_kw"]
                    if kw is None:
                        target = _worst_service(obs)
                    else:
                        target = next(
                            (s.name for s in obs.services if kw in s.name.lower()),
                            _worst_service(obs),
                        )
                    if target:
                        return Action(action_type=action_type, service=target)
                elif "reason" in kwargs:
                    return Action(action_type=ActionType.ALERT_ONCALL, reason=kwargs["reason"])
                elif "ip_range" in kwargs:
                    return Action(action_type=ActionType.BLOCK_IP_RANGE, ip_range=kwargs["ip_range"])
                elif action_type == ActionType.CREATE_INDEX:
                    db_svc = next(
                        (s.name for s in obs.services
                         if any(t in s.name for t in ["db", "database", "postgres", "vitess"])),
                        obs.services[0].name if obs.services else None,
                    )
                    return Action(action_type=ActionType.CREATE_INDEX, service=db_svc)
                elif action_type == ActionType.FAILOVER:
                    worst = _worst_service(obs)
                    return Action(action_type=ActionType.FAILOVER, service=worst, target_region="us-west-2")

        # Fallback
        return Action(
            action_type=ActionType.ALERT_ONCALL,
            reason="unresolved incident — escalating to on-call team",
        )


# --- Benchmark Runner ---------------------------------------------------------

def benchmark(
    agent_name: str,
    agent_class,
    tasks: List[str],
    num_seeds: int,
) -> Dict[str, Dict[str, float]]:
    """
    Returns {task_id: {"mean": float, "std": float, "resolved_rate": float}}.
    Scores are from grade_episode(), guaranteed [0.0, 1.0].
    A fresh agent instance is created per task to avoid cross-task state leakage.
    """
    results = {}
    for task_id in tasks:
        scores = []
        resolved = []
        # Fresh agent per task
        agent = agent_class()
        for seed in range(num_seeds):
            # RandomAgent needs the seed to vary its RNG
            if isinstance(agent, RandomAgent):
                agent = RandomAgent(seed=seed)
            score = run_episode(task_id, seed, agent)
            scores.append(score)
            # Also track resolution rate
            env = DevOpsIncidentEnv(task_id=task_id, seed=seed)
            env.reset(seed=seed)
            # (resolution already captured inside run_episode via env.state())
            # We re-derive it from score heuristic: score > 0.5 usually = resolved
            # But better to use actual flag — replay one episode:
            # (Already computed in run_episode; we just need a flag)
            # Simple: reuse the score threshold won't be reliable; track in run_episode
            resolved.append(score)

        mean  = round(statistics.mean(scores), 4)
        std   = round(statistics.stdev(scores) if len(scores) > 1 else 0.0, 4)
        results[task_id] = {
            "mean": mean,
            "std":  std,
            "min":  round(min(scores), 4),
            "max":  round(max(scores), 4),
        }
        print(f"  {agent_name:<26} {task_id:<18} mean={mean:.4f}  std={std:.4f}  "
              f"min={results[task_id]['min']:.4f}  max={results[task_id]['max']:.4f}")
    return results


# --- Pretty Table -------------------------------------------------------------

def print_table(all_results: Dict[str, Dict[str, Dict]]):
    agents = list(all_results.keys())
    tasks  = list(next(iter(all_results.values())).keys())

    col_w = 26
    hdr_w = 20

    sep  = "=" * (hdr_w + col_w * len(agents) + 2)
    dash = "-" * (hdr_w + col_w * len(agents) + 2)

    print()
    print(sep)
    print(f"  ARIA Baseline Benchmark  (grade_episode scores, [0.0, 1.0])")
    print(f"  {NUM_SEEDS} seeds per task  |  {len(tasks)} tasks  |  {len(agents)} agents")
    print(sep)
    print(f"{'Task':<{hdr_w}}", end="")
    for a in agents:
        print(f"{a:^{col_w}}", end="")
    print()
    print(f"{'':>{hdr_w}}", end="")
    for _ in agents:
        print(f"{'mean +/- std  [min, max]':^{col_w}}", end="")
    print()
    print(dash)

    for task in tasks:
        print(f"{task:<{hdr_w}}", end="")
        for agent in agents:
            r = all_results[agent][task]
            cell = f"{r['mean']:.3f}+/-{r['std']:.3f} [{r['min']:.2f},{r['max']:.2f}]"
            print(f"{cell:^{col_w}}", end="")
        print()

    print(dash)
    print(f"{'MEAN (all tasks)':<{hdr_w}}", end="")
    for agent in agents:
        vals = [all_results[agent][t]["mean"] for t in tasks]
        cell = f"{statistics.mean(vals):.3f} avg"
        print(f"{cell:^{col_w}}", end="")
    print()
    print(sep)
    print()


# --- Main ---------------------------------------------------------------------

def main():
    print()
    print("=" * 70)
    print("  ARIA Baseline Benchmark")
    print(f"  Scoring: grade_episode() — guaranteed [0.0, 1.0]")
    print(f"  Tasks: {len(ALL_TASKS)}  |  Seeds per task: {NUM_SEEDS}  |  Agents: 3")
    print("=" * 70)

    agents_config = [
        ("RandomAgent",           RandomAgent),
        ("RuleBasedAgent",        RuleBasedAgent),
        ("InformationFirstAgent", InformationFirstAgent),
    ]

    all_results = {}
    for agent_name, agent_class in agents_config:
        print(f"\n--- Running: {agent_name} ---")
        all_results[agent_name] = benchmark(agent_name, agent_class, ALL_TASKS, NUM_SEEDS)

    print_table(all_results)

    # Sanity checks
    print("=== Sanity checks ===")
    any_nonzero_std = False
    for agent in all_results:
        for task in all_results[agent]:
            r = all_results[agent][task]
            assert 0.0 <= r["mean"] <= 1.0, f"MEAN OUT OF RANGE: {agent}/{task} mean={r['mean']}"
            assert 0.0 <= r["min"] <= 1.0,  f"MIN OUT OF RANGE: {agent}/{task} min={r['min']}"
            assert 0.0 <= r["max"] <= 1.0,  f"MAX OUT OF RANGE: {agent}/{task} max={r['max']}"
            if r["std"] > 0.0:
                any_nonzero_std = True
    assert any_nonzero_std, "All std deviations are 0.000 — seed variation not working!"
    print(f"  All scores in [0.0, 1.0]: PASSED")
    print(f"  Non-zero std deviation present: PASSED")

    # Save to JSON
    output = {
        "metadata": {
            "scoring": "grade_episode() — [0.0, 1.0]",
            "num_seeds": NUM_SEEDS,
            "tasks": ALL_TASKS,
            "agents": [name for name, _ in agents_config],
            "generated_at": datetime.now(timezone.utc).isoformat(),
        },
        "results": all_results,
    }
    with open("benchmark_results.json", "w") as f:
        json.dump(output, f, indent=2)
    print("Results saved to benchmark_results.json")


if __name__ == "__main__":
    main()