"""Behavioral fingerprinting for salami-slicing attack detection.

Tracks per-agent behavioral baselines and detects gradual drift using
a two-window approach:

  • Baseline window  (first BASELINE_WINDOW actions) — establishes normal
  • Rolling window   (last  ROLLING_WINDOW  actions) — tracks current state

Salami-slicing is confirmed when the rolling average risk score exceeds
the established baseline by ≥ SALAMI_DELTA for ≥ SALAMI_SUSTAINED
consecutive scored actions.  No single action needs to cross the block
threshold — the attack is detected through cumulative drift.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional

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

BASELINE_WINDOW  = 20   # actions needed to lock the baseline (must be ≥ 1)
ROLLING_WINDOW   = 30   # rolling window for current behaviour comparison
SALAMI_DELTA     = 15   # risk-score delta above baseline to start counting
SALAMI_SUSTAINED =  8   # consecutive above-delta readings to confirm attack


# ---------------------------------------------------------------------------
# Result types
# ---------------------------------------------------------------------------

@dataclass
class DriftVector:
    """Explains *why* drift was detected — used for human-readable alerts."""

    risk_elevation: float = 0.0
    """How many points the rolling avg is above the baseline."""

    top_tool_shift: Optional[str] = None
    """The tool called most frequently in the suspicious window."""

    attack_type_surge: Optional[str] = None
    """The attack type with the highest count in the rolling window."""

    consecutive_above_threshold: int = 0
    """How many consecutive actions have been above the SALAMI_DELTA mark."""

    def to_dict(self) -> dict:
        return {
            "risk_elevation":               self.risk_elevation,
            "top_tool_shift":               self.top_tool_shift,
            "attack_type_surge":            self.attack_type_surge,
            "consecutive_above_threshold":  self.consecutive_above_threshold,
        }


@dataclass
class DriftReport:
    """Snapshot of drift metrics returned after each :meth:`update` call."""

    drift_score: int
    """0-100.  0 = no change from baseline; 100 = maximum observed drift."""

    cumulative_risk: float
    """Current rolling-window average risk score."""

    baseline_risk: float
    """Established baseline average risk score."""

    salami_detected: bool
    """True once sustained elevation has been confirmed (sticky — never reverts)."""

    drift_vector: DriftVector
    """Detailed breakdown of what is changing."""

    trend: str
    """One of: "stable" | "drifting" | "accelerating" | "compromised"."""


# ---------------------------------------------------------------------------
# BehavioralFingerprint
# ---------------------------------------------------------------------------

class BehavioralFingerprint:
    """Maintains a live behavioral fingerprint for a single agent.

    Create **one instance per** ``ProtectedAgent``; call :meth:`update` after
    every scored action to keep the fingerprint current.

    Lifecycle
    ---------
    1. **Baseline phase** — first ``BASELINE_WINDOW`` actions build the
       agent's normal risk profile.
    2. **Active monitoring** — subsequent actions are compared against the
       locked baseline via a rolling window.
    3. **Salami detection** — sustained elevation above the baseline is
       flagged as a salami-slicing attack.
    """

    def __init__(self, agent_id: str, org_id: str) -> None:
        self.agent_id = agent_id
        self.org_id   = org_id

        # --- Baseline (built during first BASELINE_WINDOW actions) ----------
        self.baseline_risk: float       = 0.0
        self.baseline_sample_count: int = 0
        self.baseline_locked: bool      = False

        # --- Rolling window of raw risk scores ------------------------------
        self._risk_window: list[int] = []

        # --- Frequency maps (explain drift-vector) --------------------------
        self.tool_frequency:     dict[str, int] = {}
        self.attack_type_counts: dict[str, int] = {}

        # --- Computed values (refreshed on every update()) ------------------
        self.current_risk: float   = 0.0
        self.drift_score: int      = 0
        self.salami_detected: bool = False
        self.drift_vector: dict    = {}

        # --- Internal salami tracker ----------------------------------------
        self._consecutive_above: int = 0

        # --- Metadata -------------------------------------------------------
        self.sessions_tracked: int       = 0
        self.total_actions: int          = 0
        self._last_session_id: Optional[str] = None

    # ------------------------------------------------------------------
    # Public interface
    # ------------------------------------------------------------------

    def update(
        self,
        risk_score: int,
        tool_name: Optional[str],
        attack_type: Optional[str],
        session_id: str,
    ) -> DriftReport:
        """Record a scored action and return the latest :class:`DriftReport`.

        This method is the **only** public mutator; call it once per
        ``score_action()`` invocation in :class:`ProtectedAgent`.

        Parameters
        ----------
        risk_score:
            The 0-100 risk score returned by the intent-scoring engine.
        tool_name:
            Name of the tool being invoked (may be None for input scans).
        attack_type:
            Attack type label from intent scoring (``"none"`` or a real type).
        session_id:
            Current session identifier — used to count distinct sessions.
        """
        self._track_session(session_id)
        self.total_actions += 1

        self._update_baseline(risk_score)
        self._update_rolling_window(risk_score)
        self._update_frequencies(tool_name, attack_type)
        self._recompute_current_risk()

        report = self._compute_drift()

        # Persist computed values for to_api_payload() serialisation.
        self.drift_score     = report.drift_score
        self.salami_detected = report.salami_detected
        self.drift_vector    = report.drift_vector.to_dict()

        return report

    def to_api_payload(self) -> dict:
        """Serialise this fingerprint for ``POST /v1/agents/{id}/fingerprint``."""
        return {
            "tool_frequency":         self.tool_frequency,
            "attack_type_counts":     self.attack_type_counts,
            "baseline_risk":          round(self.baseline_risk, 4),
            "baseline_sample_count":  self.baseline_sample_count,
            "baseline_locked":        self.baseline_locked,
            "current_risk":           round(self.current_risk, 4),
            "drift_score":            self.drift_score,
            "salami_detected":        self.salami_detected,
            "drift_vector":           self.drift_vector,
            "sessions_tracked":       self.sessions_tracked,
            "total_actions":          self.total_actions,
        }

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _track_session(self, session_id: str) -> None:
        if session_id != self._last_session_id:
            self.sessions_tracked += 1
            self._last_session_id = session_id

    def _update_baseline(self, risk_score: int) -> None:
        if self.baseline_locked:
            return
        # Running average using Welford-style one-pass formula.
        n = self.baseline_sample_count
        self.baseline_risk = (self.baseline_risk * n + risk_score) / (n + 1)
        self.baseline_sample_count += 1
        if self.baseline_sample_count >= BASELINE_WINDOW:
            self.baseline_locked = True

    def _update_rolling_window(self, risk_score: int) -> None:
        self._risk_window.append(risk_score)
        if len(self._risk_window) > ROLLING_WINDOW:
            self._risk_window.pop(0)

    def _update_frequencies(
        self,
        tool_name: Optional[str],
        attack_type: Optional[str],
    ) -> None:
        if tool_name:
            self.tool_frequency[tool_name] = (
                self.tool_frequency.get(tool_name, 0) + 1
            )
        if attack_type and attack_type != "none":
            self.attack_type_counts[attack_type] = (
                self.attack_type_counts.get(attack_type, 0) + 1
            )

    def _recompute_current_risk(self) -> None:
        if self._risk_window:
            self.current_risk = sum(self._risk_window) / len(self._risk_window)

    def _compute_drift(self) -> DriftReport:
        """Core drift computation — only meaningful once baseline is locked."""
        if not self.baseline_locked or self.baseline_risk == 0:
            # Not enough history yet; return a neutral zero-drift report.
            return DriftReport(
                drift_score=0,
                cumulative_risk=round(self.current_risk, 2),
                baseline_risk=round(self.baseline_risk, 2),
                salami_detected=False,
                drift_vector=DriftVector(),
                trend="stable",
            )

        delta = self.current_risk - self.baseline_risk

        # Drift score: linear mapping delta=0→0, delta=50→100 (clamped).
        drift_score = min(100, max(0, round(delta * 2.0)))

        # Salami counter: increment while above threshold, decay otherwise.
        if delta >= SALAMI_DELTA:
            self._consecutive_above += 1
        else:
            # Gradual decay allows brief dips without resetting the counter.
            self._consecutive_above = max(0, self._consecutive_above - 1)

        # Sticky: once confirmed, salami_detected never reverts to False.
        salami_detected = (
            self._consecutive_above >= SALAMI_SUSTAINED
            or self.salami_detected
        )

        # Drift vector — explains WHAT is changing.
        top_tool   = _top_key(self.tool_frequency)
        top_attack = _top_key(self.attack_type_counts)
        drift_vector = DriftVector(
            risk_elevation=round(delta, 2),
            top_tool_shift=top_tool,
            attack_type_surge=top_attack if self.attack_type_counts else None,
            consecutive_above_threshold=self._consecutive_above,
        )

        # Trend label.
        if salami_detected or drift_score >= 70:
            trend = "compromised"
        elif drift_score >= 40:
            trend = "accelerating"
        elif drift_score >= 15:
            trend = "drifting"
        else:
            trend = "stable"

        return DriftReport(
            drift_score=drift_score,
            cumulative_risk=round(self.current_risk, 2),
            baseline_risk=round(self.baseline_risk, 2),
            salami_detected=salami_detected,
            drift_vector=drift_vector,
            trend=trend,
        )


# ---------------------------------------------------------------------------
# Utility
# ---------------------------------------------------------------------------

def _top_key(d: dict) -> Optional[str]:
    """Return the key with the highest integer value, or None if empty."""
    return max(d, key=d.__getitem__) if d else None
