# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
"""Rogue agent detection engine (OWASP ASI-10).

Detects behavioral anomalies that indicate an agent has gone rogue:
tool-call frequency spikes, action entropy deviation, and capability
profile violations.  Combines these signals into a single risk
assessment with optional auto-quarantine recommendations.
"""

from __future__ import annotations

import logging
import math
import time
from collections import defaultdict, deque
from dataclasses import dataclass, field
from enum import Enum
from typing import Any


logger = logging.getLogger(__name__)


# ── Enums / value types ─────────────────────────────────────────────


class RiskLevel(Enum):
    """Overall risk classification for an agent."""

    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"


# ── Configuration ────────────────────────────────────────────────────


@dataclass
class RogueDetectorConfig:
    """Tunable thresholds for the rogue-agent detector."""

    # ToolCallFrequencyAnalyzer
    frequency_window_seconds: float = 60.0
    frequency_z_threshold: float = 2.5
    frequency_min_windows: int = 5

    # ActionEntropyScorer
    entropy_low_threshold: float = 0.3
    entropy_high_threshold: float = 3.5
    entropy_min_actions: int = 10

    # CapabilityProfileDeviation
    capability_violation_weight: float = 1.0

    # RogueAgentDetector
    quarantine_risk_level: RiskLevel = RiskLevel.HIGH


# ── Assessment result ────────────────────────────────────────────────


@dataclass
class RogueAssessment:
    """Result of a rogue-agent risk assessment."""

    agent_id: str
    risk_level: RiskLevel
    composite_score: float
    frequency_score: float
    entropy_score: float
    capability_score: float
    quarantine_recommended: bool
    details: dict[str, Any] = field(default_factory=dict)
    timestamp: float = field(default_factory=time.time)

    def to_dict(self) -> dict[str, Any]:
        return {
            "agent_id": self.agent_id,
            "risk_level": self.risk_level.value,
            "composite_score": round(self.composite_score, 4),
            "frequency_score": round(self.frequency_score, 4),
            "entropy_score": round(self.entropy_score, 4),
            "capability_score": round(self.capability_score, 4),
            "quarantine_recommended": self.quarantine_recommended,
            "details": self.details,
            "timestamp": self.timestamp,
        }


# ── Analyzers ────────────────────────────────────────────────────────


class ToolCallFrequencyAnalyzer:
    """Z-score analysis on tool-call frequency per time window.

    Tracks calls per agent per time bucket and detects sudden spikes
    using z-score against baseline.
    """

    def __init__(
        self,
        window_seconds: float = 60.0,
        z_threshold: float = 2.5,
        min_windows: int = 5,
    ) -> None:
        self.window_seconds = window_seconds
        self.z_threshold = z_threshold
        self.min_windows = min_windows
        # agent_id → deque of (window_start, count) tuples
        self._buckets: dict[str, deque[tuple[float, int]]] = defaultdict(
            lambda: deque(maxlen=200),
        )
        # agent_id → (current_window_start, running_count)
        self._current: dict[str, tuple[float, int]] = {}

    def _flush_bucket(self, agent_id: str, now: float) -> None:
        """Rotate the current bucket if the window has elapsed."""
        if agent_id not in self._current:
            self._current[agent_id] = (now, 0)
            return

        win_start, count = self._current[agent_id]
        if now - win_start >= self.window_seconds:
            self._buckets[agent_id].append((win_start, count))
            self._current[agent_id] = (now, 0)

    def record(self, agent_id: str, timestamp: float | None = None) -> None:
        """Record a single tool call."""
        now = timestamp if timestamp is not None else time.time()
        self._flush_bucket(agent_id, now)
        win_start, count = self._current[agent_id]
        self._current[agent_id] = (win_start, count + 1)

    def score(self, agent_id: str, timestamp: float | None = None) -> float:
        """Return a z-score for the current window's call frequency.

        Returns 0.0 when insufficient data is available.
        """
        now = timestamp if timestamp is not None else time.time()
        self._flush_bucket(agent_id, now)

        history = self._buckets.get(agent_id)
        if not history or len(history) < self.min_windows:
            return 0.0

        counts = [c for _, c in history]
        mean = sum(counts) / len(counts)
        variance = sum((x - mean) ** 2 for x in counts) / len(counts)
        std_dev = math.sqrt(variance)

        _, current_count = self._current.get(agent_id, (0.0, 0))

        if std_dev == 0:
            return 0.0

        return abs(current_count - mean) / std_dev


class ActionEntropyScorer:
    """Shannon entropy scoring for agent action diversity.

    Low entropy → agent stuck in a loop (repetitive behavior).
    Very high entropy → erratic, unfocused behavior.
    """

    def __init__(
        self,
        low_threshold: float = 0.3,
        high_threshold: float = 3.5,
        min_actions: int = 10,
    ) -> None:
        self.low_threshold = low_threshold
        self.high_threshold = high_threshold
        self.min_actions = min_actions
        # agent_id → list of action names
        self._actions: dict[str, list[str]] = defaultdict(list)

    def record(self, agent_id: str, action: str) -> None:
        """Record an action for the given agent."""
        self._actions[agent_id].append(action)

    def entropy(self, agent_id: str) -> float | None:
        """Compute Shannon entropy of the agent's action distribution.

        Returns ``None`` when fewer than ``min_actions`` have been recorded.
        """
        actions = self._actions.get(agent_id, [])
        if len(actions) < self.min_actions:
            return None

        total = len(actions)
        counts: dict[str, int] = {}
        for a in actions:
            counts[a] = counts.get(a, 0) + 1

        h = 0.0
        for c in counts.values():
            p = c / total
            if p > 0:
                h -= p * math.log2(p)
        return h

    def score(self, agent_id: str) -> float:
        """Return an anomaly score based on entropy deviation.

        * Score > 0 indicates anomalous entropy (too low or too high).
        * Returns 0.0 when data is insufficient or entropy is in range.
        """
        h = self.entropy(agent_id)
        if h is None:
            return 0.0

        if h < self.low_threshold:
            # Agent stuck in a loop — further below threshold → higher score
            return (self.low_threshold - h) / self.low_threshold if self.low_threshold > 0 else 1.0

        if h > self.high_threshold:
            # Erratic behavior
            return (h - self.high_threshold) / self.high_threshold if self.high_threshold > 0 else 1.0

        return 0.0


class CapabilityProfileDeviation:
    """Detect when an agent uses tools outside its declared profile.

    Each agent registers a set of allowed tools.  Any tool call
    outside that set increments a violation counter.  The score
    is the fraction of total calls that are violations.
    """

    def __init__(self, violation_weight: float = 1.0) -> None:
        self.violation_weight = violation_weight
        # agent_id → set of allowed tool names
        self._profiles: dict[str, set[str]] = {}
        # agent_id → (total_calls, violation_count)
        self._counters: dict[str, tuple[int, int]] = defaultdict(lambda: (0, 0))

    def register_profile(self, agent_id: str, allowed_tools: list[str]) -> None:
        """Register or update the capability profile for an agent."""
        self._profiles[agent_id] = set(allowed_tools)

    def record(self, agent_id: str, tool_name: str) -> bool:
        """Record a tool call and return ``True`` if it is a violation."""
        total, violations = self._counters[agent_id]
        total += 1

        is_violation = False
        profile = self._profiles.get(agent_id)
        if profile is not None and tool_name not in profile:
            violations += 1
            is_violation = True

        self._counters[agent_id] = (total, violations)
        return is_violation

    def score(self, agent_id: str) -> float:
        """Return a capability-deviation score in ``[0, 1]``.

        A value of 0 means no violations; 1 means every call was
        outside the declared profile.  Scaled by ``violation_weight``.
        """
        total, violations = self._counters.get(agent_id, (0, 0))
        if total == 0:
            return 0.0
        return (violations / total) * self.violation_weight


# ── Orchestrator ─────────────────────────────────────────────────────

_RISK_ORDER = [RiskLevel.LOW, RiskLevel.MEDIUM, RiskLevel.HIGH, RiskLevel.CRITICAL]


class RogueAgentDetector:
    """Orchestrates frequency, entropy, and capability analyzers to
    produce a composite risk assessment for each agent.
    """

    def __init__(self, config: RogueDetectorConfig | None = None) -> None:
        self._config = config or RogueDetectorConfig()

        self.frequency_analyzer = ToolCallFrequencyAnalyzer(
            window_seconds=self._config.frequency_window_seconds,
            z_threshold=self._config.frequency_z_threshold,
            min_windows=self._config.frequency_min_windows,
        )
        self.entropy_scorer = ActionEntropyScorer(
            low_threshold=self._config.entropy_low_threshold,
            high_threshold=self._config.entropy_high_threshold,
            min_actions=self._config.entropy_min_actions,
        )
        self.capability_checker = CapabilityProfileDeviation(
            violation_weight=self._config.capability_violation_weight,
        )

        self._assessments: list[RogueAssessment] = []

    # -- data ingestion ---------------------------------------------------

    def record_action(
        self,
        agent_id: str,
        action: str,
        tool_name: str,
        timestamp: float | None = None,
    ) -> None:
        """Feed an observed action into all analyzers."""
        ts = timestamp if timestamp is not None else time.time()
        self.frequency_analyzer.record(agent_id, timestamp=ts)
        self.entropy_scorer.record(agent_id, action)
        self.capability_checker.record(agent_id, tool_name)

    def register_capability_profile(
        self,
        agent_id: str,
        allowed_tools: list[str],
    ) -> None:
        """Declare the set of tools an agent is expected to use."""
        self.capability_checker.register_profile(agent_id, allowed_tools)

    # -- assessment -------------------------------------------------------

    def assess(
        self,
        agent_id: str,
        timestamp: float | None = None,
    ) -> RogueAssessment:
        """Produce a composite risk assessment for *agent_id*.

        Combines frequency z-score, entropy anomaly score, and
        capability violation score into a single ``RogueAssessment``.
        """
        freq_score = self.frequency_analyzer.score(agent_id, timestamp=timestamp)
        ent_score = self.entropy_scorer.score(agent_id)
        cap_score = self.capability_checker.score(agent_id)

        composite = freq_score + ent_score + cap_score
        risk_level = self._classify_risk(composite)

        quarantine_threshold_idx = _RISK_ORDER.index(
            self._config.quarantine_risk_level,
        )
        current_idx = _RISK_ORDER.index(risk_level)
        quarantine = current_idx >= quarantine_threshold_idx

        assessment = RogueAssessment(
            agent_id=agent_id,
            risk_level=risk_level,
            composite_score=composite,
            frequency_score=freq_score,
            entropy_score=ent_score,
            capability_score=cap_score,
            quarantine_recommended=quarantine,
            details={
                "frequency_z_threshold": self._config.frequency_z_threshold,
                "entropy_low_threshold": self._config.entropy_low_threshold,
                "entropy_high_threshold": self._config.entropy_high_threshold,
                "quarantine_risk_level": self._config.quarantine_risk_level.value,
            },
            timestamp=timestamp if timestamp is not None else time.time(),
        )

        self._assessments.append(assessment)
        if quarantine:
            logger.warning(
                "Quarantine recommended for agent %s (risk=%s, score=%.2f)",
                agent_id,
                risk_level.value,
                composite,
            )

        return assessment

    # -- queries ----------------------------------------------------------

    @property
    def assessments(self) -> list[RogueAssessment]:
        """Return a copy of the assessment history."""
        return list(self._assessments)

    # -- internals --------------------------------------------------------

    @staticmethod
    def _classify_risk(composite_score: float) -> RiskLevel:
        """Map composite score to a risk level.

        Thresholds (mirroring ``AnomalySeverity`` boundaries):
        * < 1.0  → LOW
        * < 2.0  → MEDIUM
        * < 3.0  → HIGH
        * >= 3.0 → CRITICAL
        """
        if composite_score >= 3.0:
            return RiskLevel.CRITICAL
        if composite_score >= 2.0:
            return RiskLevel.HIGH
        if composite_score >= 1.0:
            return RiskLevel.MEDIUM
        return RiskLevel.LOW
