"""Local behavioral baseline — per-agent deviation detection on-device.

Runs entirely on the user's machine. Zero external dependencies (stdlib only).
Learns what "normal" looks like for an agent over time, then flags deviations.

No content leaves the device. Only deviation flag names are used in telemetry.

Tracked metrics:
  - action_distribution: counter of action types (Read, Write, Bash, etc.)
  - action_entropy: Shannon entropy of action distribution
  - hourly_distribution: counter of events per hour-of-day (0-23 UTC)
  - session_lengths: rolling window of recent session event counts
  - error_rate: fraction of events with error result
  - total_events: total events seen

Deviation types (requires >= 100 events):
  - many_novel_actions: >30% of session actions never seen in baseline
  - entropy_shift: entropy diff >0.5 from baseline
  - error_rate_spike: error rate >20% above baseline
  - unusual_timing: activity outside typical hours (<2% historical activity)
  - session_length_anomaly: session >3x average session length
  - action_burst: >50 actions in 60 seconds
"""
import datetime
import hashlib
import json
import math
import os
import time


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

_MIN_EVENTS = 100
_SESSION_WINDOW = 50       # rolling window for session_lengths
_NOVEL_ACTION_THRESHOLD = 0.3
_ENTROPY_THRESHOLD = 0.5
_ERROR_RATE_THRESHOLD = 0.20
_TIMING_THRESHOLD = 0.02   # less than 2% of historical activity
_SESSION_LENGTH_MULTIPLIER = 3
_BURST_COUNT = 50
_BURST_WINDOW = 60.0       # seconds
_BASELINE_VERSION = 1
_SAVE_INTERVAL = 10        # persist to disk every N events

# Actions that should NEVER be normalized into baseline.
# These are inherently dangerous and an attacker could use gradual drift
# to shift the baseline to include them, suppressing deviation alerts.
_BASELINE_EXCLUSIONS = frozenset({
    # Detection flag names that should NEVER be normalized into baseline.
    # An attacker could gradually drift the baseline by introducing these
    # actions slowly; we prevent this by never learning from flagged events.
    "reverse_shell", "crypto_miner", "container_escape",
    "process_injection", "kernel_module_loading",
    "escape_attempt", "privilege_escalation",
    "piped_exfiltration", "delayed_exfiltration",
    "variable_exfiltration", "data_staging",
    "persistence_install", "agent_self_modification",
    "data_destruction", "alias_obfuscation",
    "symlink_indirection", "supply_chain_risk",
    "dns_exfiltration", "env_harvesting",
})


# ---- Entropy -----------------------------------------------------------------

def _shannon_entropy(counts: dict) -> float:
    """Compute Shannon entropy (bits) from a counter dict."""
    total = sum(counts.values())
    if total == 0:
        return 0.0
    probs = [c / total for c in counts.values() if c > 0]
    return -sum(p * math.log2(p) for p in probs)


# ---- LocalBaseline -----------------------------------------------------------

class LocalBaseline:
    """Per-agent behavioral baseline that learns and detects deviations.

    Stores baseline data at ``~/.obol/baselines/{agent_id_hash}.json``.
    Atomic writes (tmp + rename) with restrictive permissions (0o600).
    Never raises exceptions -- returns empty list on any failure.
    """

    def __init__(self, agent_id: str, baseline_dir: str = None):
        """Load or create baseline for *agent_id*.

        Parameters
        ----------
        agent_id : str
            Unique identifier for the agent being monitored.
        baseline_dir : str, optional
            Directory for baseline files.  Defaults to ``~/.obol/baselines/``.
        """
        try:
            self._agent_id = agent_id
            self._id_hash = hashlib.sha256(agent_id.encode()).hexdigest()

            if baseline_dir is None:
                baseline_dir = os.path.join(
                    os.path.expanduser("~"), ".obol", "baselines",
                )
            self._baseline_dir = baseline_dir
            self._path = os.path.join(self._baseline_dir, f"{self._id_hash}.json")

            # Recent event timestamps for burst detection (in-memory only)
            self._recent_timestamps: list = []

            self._data = self._load()
        except Exception:
            self._data = self._empty()
            self._recent_timestamps = []

    # ---- Public API -----------------------------------------------------------

    def check_and_update(self, action: str, result: str = "success",
                         session_event_count: int = 0,
                         session_actions: dict = None,
                         detected_flags: list = None) -> list:
        """Check for deviations from baseline, then update baseline.

        Parameters
        ----------
        action : str
            Tool action name (e.g. "Read", "Write", "Bash").
        result : str
            Outcome -- ``"success"`` or ``"error"``.
        session_event_count : int
            How many events in the current session so far.
        session_actions : dict
            Mapping of ``{action: count}`` for the current session.
        detected_flags : list, optional
            Detection flags from the current event. If any match
            ``_BASELINE_EXCLUSIONS``, the baseline is NOT updated
            (prevents gradual drift attacks from normalizing malicious
            actions into the baseline).

        Returns
        -------
        list[str]
            Deviation flag strings (empty if normal).
        """
        try:
            # 1. Check deviations BEFORE updating (compare against prior state)
            flags = self._check_deviations(
                action, result, session_event_count, session_actions,
            )

            # 2. Update baseline — UNLESS event has excluded detection flags.
            #    An attacker can gradually drift the baseline by introducing
            #    dangerous actions slowly. We prevent this by never learning
            #    from events that triggered high-risk detection flags.
            skip_update = False
            if detected_flags:
                for flag in detected_flags:
                    if flag in _BASELINE_EXCLUSIONS:
                        skip_update = True
                        break

            if not skip_update:
                self._update(action, result, session_event_count)

            # 3. Persist (every N events or when deviations detected)
            total = self._data.get("total_events", 0)
            if flags or total % _SAVE_INTERVAL == 0:
                self.save()

            return flags
        except Exception:
            return []

    def save(self):
        """Persist baseline to disk.  Called automatically by check_and_update."""
        try:
            os.makedirs(self._baseline_dir, mode=0o700, exist_ok=True)
            # Reject world-writable baseline directory (symlink attack surface)
            dir_stat = os.stat(self._baseline_dir)
            if dir_stat.st_mode & 0o002:  # world-writable
                import sys
                print(
                    f"[obol] baseline dir {self._baseline_dir} is world-writable, "
                    "refusing to save (security risk)",
                    file=sys.stderr,
                )
                return
            tmp = self._path + ".tmp"
            fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
            with os.fdopen(fd, "w") as f:
                json.dump(self._data, f)
            os.replace(tmp, self._path)
        except Exception:
            pass

    @property
    def is_ready(self) -> bool:
        """True if enough events to detect deviations (>= 100)."""
        try:
            return self._data.get("total_events", 0) >= _MIN_EVENTS
        except Exception:
            return False

    @property
    def summary(self) -> dict:
        """Return summary stats for telemetry (no PII, just shapes).

        Keys: total_events, action_count, entropy, avg_session_length,
              error_rate, typical_hours, is_ready.
        """
        try:
            d = self._data
            action_dist = d.get("action_distribution", {})
            session_lengths = d.get("session_lengths", [])
            hourly = d.get("hourly_distribution", {})
            total_hourly = sum(hourly.values()) if hourly else 0

            avg_session = 0.0
            if session_lengths:
                avg_session = sum(session_lengths) / len(session_lengths)

            # Typical hours: those with >= 5% of historical activity
            typical = []
            if total_hourly > 0:
                typical = sorted(
                    int(h) for h, c in hourly.items()
                    if c / total_hourly >= 0.05
                )

            return {
                "total_events": d.get("total_events", 0),
                "action_count": len(action_dist),
                "entropy": round(_shannon_entropy(action_dist), 4),
                "avg_session_length": round(avg_session, 2),
                "error_rate": round(d.get("error_rate", 0.0), 4),
                "typical_hours": typical,
                "is_ready": self.is_ready,
            }
        except Exception:
            return {
                "total_events": 0,
                "action_count": 0,
                "entropy": 0.0,
                "avg_session_length": 0.0,
                "error_rate": 0.0,
                "typical_hours": [],
                "is_ready": False,
            }

    # ---- Internal -------------------------------------------------------------

    def _empty(self) -> dict:
        """Return a fresh, empty baseline structure."""
        return {
            "action_distribution": {},
            "hourly_distribution": {},
            "session_lengths": [],
            "error_rate": 0.0,
            "total_events": 0,
            "last_updated": None,
            "baseline_version": _BASELINE_VERSION,
        }

    def _load(self) -> dict:
        """Load baseline from disk, or return empty if missing/corrupt."""
        try:
            with open(self._path, "r") as f:
                raw = f.read(2 * 1024 * 1024)  # 2 MB cap
            data = json.loads(raw)
            if not isinstance(data, dict):
                return self._empty()
            # Ensure required keys exist (forward compat)
            empty = self._empty()
            for key in empty:
                if key not in data:
                    data[key] = empty[key]
            # Validate expected types — reset to defaults if corrupted
            data = self._validate(data)
            return data
        except (FileNotFoundError, json.JSONDecodeError, OSError, ValueError):
            return self._empty()

    @staticmethod
    def _validate(data: dict) -> dict:
        """Validate and sanitize loaded baseline data.

        Resets individual fields to safe defaults when types are wrong,
        and bounds collection sizes to prevent unbounded growth from
        corrupt or tampered files.
        """
        # action_distribution must be dict[str, int/float]
        ad = data.get("action_distribution")
        if not isinstance(ad, dict):
            data["action_distribution"] = {}
        else:
            # Bound to 500 unique actions (prevent memory abuse)
            cleaned = {}
            for k, v in list(ad.items())[:500]:
                if isinstance(k, str) and isinstance(v, (int, float)) and v >= 0:
                    cleaned[str(k)] = v
            data["action_distribution"] = cleaned

        # hourly_distribution must be dict[str(0-23), int/float]
        hd = data.get("hourly_distribution")
        if not isinstance(hd, dict):
            data["hourly_distribution"] = {}
        else:
            cleaned = {}
            for k, v in hd.items():
                if isinstance(v, (int, float)) and v >= 0:
                    cleaned[str(k)] = v
            data["hourly_distribution"] = cleaned

        # session_lengths must be list[int/float], bounded
        sl = data.get("session_lengths")
        if not isinstance(sl, list):
            data["session_lengths"] = []
        else:
            data["session_lengths"] = [
                x for x in sl[-_SESSION_WINDOW:]
                if isinstance(x, (int, float)) and x >= 0
            ]

        # error_rate must be float in [0.0, 1.0]
        er = data.get("error_rate")
        if not isinstance(er, (int, float)) or er < 0 or er > 1:
            data["error_rate"] = 0.0
        else:
            data["error_rate"] = float(er)

        # total_events must be non-negative int
        te = data.get("total_events")
        if not isinstance(te, (int, float)) or te < 0:
            data["total_events"] = 0
        else:
            data["total_events"] = int(te)

        # baseline_version must be int
        bv = data.get("baseline_version")
        if not isinstance(bv, int):
            data["baseline_version"] = _BASELINE_VERSION

        return data

    def _update(self, action: str, result: str, session_event_count: int):
        """Incorporate current event into baseline."""
        d = self._data

        # Action distribution
        ad = d.setdefault("action_distribution", {})
        ad[action] = ad.get(action, 0) + 1

        # Hourly distribution
        hour = str(datetime.datetime.now(datetime.timezone.utc).hour)
        hd = d.setdefault("hourly_distribution", {})
        hd[hour] = hd.get(hour, 0) + 1

        # Total events and error rate (incremental)
        total = d.get("total_events", 0)
        old_rate = d.get("error_rate", 0.0)
        total += 1
        if result == "error":
            d["error_rate"] = old_rate + (1.0 - old_rate) / total
        else:
            d["error_rate"] = old_rate - (old_rate / total)
        d["total_events"] = total

        # Session lengths (append when session_event_count > 0 and looks like
        # a session boundary, i.e. count is smaller than previous observation).
        # We record the session length once when we see a new session start.
        sl = d.setdefault("session_lengths", [])
        if session_event_count > 0:
            # Record every incoming count; the caller increments per-event.
            # We overwrite the last entry if it belongs to the same session
            # (strictly increasing counts), and start a new entry otherwise.
            if sl and session_event_count > sl[-1]:
                sl[-1] = session_event_count
            else:
                sl.append(session_event_count)
            # Rolling window
            if len(sl) > _SESSION_WINDOW:
                d["session_lengths"] = sl[-_SESSION_WINDOW:]

        # Burst tracking (in-memory)
        now = time.monotonic()
        self._recent_timestamps.append(now)
        cutoff = now - _BURST_WINDOW
        self._recent_timestamps = [
            t for t in self._recent_timestamps if t >= cutoff
        ]

        d["last_updated"] = (
            datetime.datetime.now(datetime.timezone.utc)
            .replace(tzinfo=None).isoformat() + "Z"
        )

    def _check_deviations(self, action: str, result: str,
                          session_event_count: int,
                          session_actions: dict) -> list:
        """Compare current event/session against baseline and return flags."""
        flags: list = []
        if not self.is_ready:
            return flags

        d = self._data

        # ---- many_novel_actions ------------------------------------------------
        if session_actions:
            known = set(d.get("action_distribution", {}).keys())
            session_set = set(session_actions.keys())
            novel = session_set - known
            if len(novel) / max(len(session_set), 1) > _NOVEL_ACTION_THRESHOLD:
                flags.append("many_novel_actions")

        # ---- entropy_shift -----------------------------------------------------
        if session_actions:
            session_entropy = _shannon_entropy(session_actions)
            baseline_entropy = _shannon_entropy(
                d.get("action_distribution", {}),
            )
            if abs(session_entropy - baseline_entropy) > _ENTROPY_THRESHOLD:
                flags.append("entropy_shift")

        # ---- error_rate_spike --------------------------------------------------
        baseline_error_rate = d.get("error_rate", 0.0)
        if session_actions and session_event_count > 10:
            # Compute actual session error rate from session_actions
            # (session_actions is {action: count}; errors tracked separately
            # in SessionDetector, so approximate from session_event_count)
            # Only flag if THIS event is an error AND session shows elevated rate
            if result == "error":
                # Estimate: if baseline error rate is <5% and we're seeing errors,
                # that's suspicious. But require more than just one error —
                # check if error_rate is significantly above baseline.
                spike_threshold = max(baseline_error_rate + _ERROR_RATE_THRESHOLD, 0.15)
                # We can't compute exact session error rate here (no error count),
                # but we CAN check if baseline is clean and this is unusual
                if baseline_error_rate < 0.05:
                    flags.append("error_rate_spike")
                # For agents with moderate error rates, only flag if baseline
                # is still quite clean (< 20%) to avoid noise
                elif baseline_error_rate < 0.20 and session_event_count > 20:
                    flags.append("error_rate_spike")

        # ---- unusual_timing ----------------------------------------------------
        hour = str(datetime.datetime.now(datetime.timezone.utc).hour)
        hourly = d.get("hourly_distribution", {})
        total_hourly = sum(hourly.values()) if hourly else 0
        if total_hourly > _MIN_EVENTS:
            hour_fraction = hourly.get(hour, 0) / total_hourly
            if hour_fraction < _TIMING_THRESHOLD:
                flags.append("unusual_timing")

        # ---- session_length_anomaly --------------------------------------------
        # Use median (not mean) to resist poisoning via many short sessions.
        # Floor: even if median is 1, still flag sessions >3x the median or >150.
        session_lengths = d.get("session_lengths", [])
        if session_lengths and session_event_count > 0:
            sorted_lens = sorted(session_lengths)
            n = len(sorted_lens)
            if n % 2 == 1:
                median_len = sorted_lens[n // 2]
            else:
                median_len = (sorted_lens[n // 2 - 1] + sorted_lens[n // 2]) / 2
            # Enforce a minimum threshold of 50 to prevent trivial poisoning
            threshold = max(median_len * _SESSION_LENGTH_MULTIPLIER, 50)
            if session_event_count > threshold:
                flags.append("session_length_anomaly")

        # ---- action_burst (in-memory timestamps) -------------------------------
        now = time.monotonic()
        cutoff = now - _BURST_WINDOW
        recent = [t for t in self._recent_timestamps if t >= cutoff]
        if len(recent) >= _BURST_COUNT:
            flags.append("action_burst")

        return flags
