#Baseline agent — heuristic rule-based agent for the DebugOps environment.

# Strategy
# --------
# 1. Parse log keywords to identify the most likely root cause.
# 2. Track own action history to avoid repeating wrong actions.
# 3. Fall back to metric-based decisions when logs are ambiguous.
# 4. Use the multi-step fix sequences documented in dynamics.py.
from __future__ import annotations
from typing import Dict, Any, List


# Known fix sequences per root cause (mirrors dynamics.py FIX_MAP)
_FIX_SEQUENCES: Dict[str, List[str]] = {
    "api_timeout":        ["scale_up",    "restart_api"],
    "db_connection_leak": ["restart_db",  "scale_up"],
    "cache_miss_storm":   ["restart_cache", "scale_up"],
    "memory_leak":        ["restart_api", "restart_db"],
}

# Log keywords → probable root cause
_LOG_SIGNALS: Dict[str, str] = {
    "timeout":      "api_timeout",
    "upstream":     "api_timeout",
    "connections":  "db_connection_leak",
    "db pool":      "db_connection_leak",
    "db":           "db_connection_leak",
    "cache miss":   "cache_miss_storm",
    "cache":        "cache_miss_storm",
    "memory":       "memory_leak",
    "oom":          "memory_leak",
    "heap":         "memory_leak",
}

# Module-level action history (reset at the start of each episode call)
_action_history: List[str] = []


def reset_history() -> None:
    """Call at the start of a new episode to clear action memory."""
    global _action_history
    _action_history = []


def act(state: Dict[str, Any]) -> str:
    """
    Choose the next action given the current observation.

    Parameters
    ----------
    state : dict with keys services, logs, metrics, time_step

    Returns
    -------
    str : one of restart_api | restart_db | restart_cache | scale_up | noop
    """
    logs_text = " ".join(state.get("logs", [])).lower()
    metrics   = state.get("metrics", {})

    # Step 1: identify probable root cause from logs
    cause = _infer_cause(logs_text)

    if cause:
        seq = _FIX_SEQUENCES[cause]
        # Determine how many steps of this sequence we've already issued
        progress = _count_progress(seq, _action_history)
        if progress < len(seq):
            action = seq[progress]
            _action_history.append(action)
            return action

    # Step 2: metric-based fallback
    if metrics.get("cpu", 0) > 80 and "scale_up" not in _action_history:
        _action_history.append("scale_up")
        return "scale_up"

    if metrics.get("error_rate", 0) > 0.5 and "restart_api" not in _action_history:
        _action_history.append("restart_api")
        return "restart_api"

    # Step 3: try any unused action (avoid noop loops) 
    for a in ["restart_api", "restart_db", "restart_cache", "scale_up"]:
        if _action_history.count(a) < 2:
            _action_history.append(a)
            return a

    _action_history.append("noop")
    return "noop"

def _infer_cause(logs_text: str) -> str | None:
    for keyword, cause in _LOG_SIGNALS.items():
        if keyword in logs_text:
            return cause
    return None


def _count_progress(seq: List[str], history: List[str]) -> int:
    # Count how many leading actions of `seq` appear in `history` in order.
    idx = 0
    for action in history:
        if idx < len(seq) and action == seq[idx]:
            idx += 1
    return idx