"""
difficulty_controller.py — Adaptive drift difficulty controller (v2).

Tracks agent belief_accuracy per drift_type over a sliding window.
Escalates difficulty when the agent consistently performs well above threshold.

Why this matters:
  Prevents the agent from overfitting to 12 fixed scenarios.
  The environment generates effectively infinite scenarios by combining drift
  types and adjusting parameters based on agent competence.

Escalation levels:
  0 → base scenarios unchanged
  1 → combine two drift types in same episode (secondary drift added 2 steps later)
  2 → noisy error messages (ambiguous 422 bodies — "validation error" without field name)
  3 → force silent semantic drift + noisy errors (hardest: no error signal at all)
"""
from __future__ import annotations

from typing import Dict, List
from collections import defaultdict, deque


ESCALATION_THRESHOLD = 0.80   # average belief_accuracy to trigger escalation
ESCALATION_WINDOW = 5         # number of recent episodes to average per drift type
EASY_DRIFT_TYPES = ["field_rename", "endpoint_version", "policy_change"]
HARD_DRIFT_TYPES = ["silent_semantic"]


class DriftDifficultyController:
    """
    Tracks agent performance and adaptively escalates drift difficulty.
    A single instance persists across all episodes for a given server process.
    """

    def __init__(self):
        # belief_accuracy history per drift_type
        self._history: Dict[str, deque] = defaultdict(
            lambda: deque(maxlen=ESCALATION_WINDOW)
        )
        self._escalation_level: int = 0   # 0=basic, 1=combined, 2=noisy, 3=silent-only

    def record(self, drift_type: str, belief_accuracy: float) -> None:
        """Call after each episode finishes to update the performance history."""
        self._history[drift_type].append(belief_accuracy)
        self._maybe_escalate()

    def _maybe_escalate(self) -> None:
        """
        Escalate if agent averages above threshold on ALL easy drift types
        for the last ESCALATION_WINDOW episodes.
        """
        if self._escalation_level >= 3:
            return
        for dt in EASY_DRIFT_TYPES:
            hist = list(self._history[dt])
            if len(hist) < ESCALATION_WINDOW:
                return  # not enough data yet
            avg = sum(hist) / len(hist)
            if avg < ESCALATION_THRESHOLD:
                return  # agent still struggling on this drift type
        self._escalation_level += 1

    def get_scenario_params(self, base_scenario: dict) -> dict:
        """
        Modify a base scenario dict based on current escalation level.

        Level 0: return base scenario unchanged
        Level 1: add secondary_drift_step 2 steps after primary (combines two drift types)
        Level 2: set noisy_errors=True (ambiguous error messages)
        Level 3: force drift_type to "silent_semantic" + noisy errors
        """
        if self._escalation_level == 0:
            return base_scenario

        modified = dict(base_scenario)

        if self._escalation_level >= 1:
            # Add a secondary drift 2 steps after the primary drift
            primary_step = base_scenario.get("drift_trigger_step", 3)
            modified["secondary_drift_step"] = primary_step + 2
            modified["secondary_drift_type"] = "policy_change"
            # Also set noisy_errors in initial_world if it exists
            if "initial_world" in modified:
                modified["initial_world"] = dict(modified["initial_world"])

        if self._escalation_level >= 2:
            # Noisy errors: 422 body becomes vague
            modified["noisy_errors"] = True
            # Merge into mutated_world if present
            if "mutated_world" in modified:
                modified["mutated_world"] = dict(modified["mutated_world"])
                modified["mutated_world"]["noisy_errors"] = True

        if self._escalation_level >= 3:
            # Force silent semantic only — override drift_type
            modified["drift_type"] = "silent_semantic"
            modified["noisy_errors"] = True
            if "mutated_world" in modified:
                modified["mutated_world"]["noisy_errors"] = True

        return modified

    def deescalate(self) -> None:
        """Manually lower escalation level (for testing or reset)."""
        if self._escalation_level > 0:
            self._escalation_level -= 1

    def reset(self) -> None:
        """Reset all state (for testing)."""
        self._history.clear()
        self._escalation_level = 0

    @property
    def level(self) -> int:
        return self._escalation_level

    def get_history_summary(self) -> Dict[str, float]:
        """Returns current average belief_accuracy per drift type (for logging)."""
        return {
            dt: (sum(hist) / len(hist) if hist else 0.0)
            for dt, hist in self._history.items()
        }