"""
InsAIts SDK - Behavioral Fingerprint Detector
===============================================
Detects: Rug Pull attacks (behavioral change after package update)
OWASP: MCP04 (Supply Chain Attacks)
CVE: CVE-2025-54136 (MCPoison), CVE-2025-6514 (mcp-remote)
Priority: #2

Algorithm:
    1. Build per-agent behavioral fingerprint over a warmup window
    2. Fingerprint captures: message length distribution, vocabulary profile,
       response time patterns, anomaly rate baseline, topic consistency
    3. After warmup, flag significant deviations from established fingerprint
    4. Supports manual snapshot/compare for pre/post-update verification

Observable indicators (from cve_signatures.py):
    - Tool behavioral fingerprint changes significantly after package update
    - Tool call patterns outside established baseline post-update
    - Semantic drift in tool responses compared to pre-update baseline
    - New tool calls appearing that were not in pre-update baseline
"""

import math
import re
import time
from typing import Dict, List, Optional, Set


class BehavioralFingerprintDetector:
    """Detects behavioral fingerprint changes in agents over time.

    Builds a statistical profile of each agent's communication patterns
    and flags significant deviations that may indicate a rug pull or
    behavioral modification.

    Detects: Rug Pull / Supply Chain Behavioral Change
    OWASP: MCP04
    """

    name = "behavioral_fingerprint"
    enabled = True

    # Minimum messages before fingerprint is considered stable
    WARMUP_WINDOW = 10
    # Number of recent messages to compare against baseline
    COMPARISON_WINDOW = 5

    def __init__(self):
        # Per-agent fingerprints: {agent_id: AgentFingerprint}
        self._fingerprints: Dict[str, Dict] = {}
        # Snapshots for pre/post-update comparison
        self._snapshots: Dict[str, Dict] = {}

    def check(self, text: str, message_idx: int,
              agent_id: str = "unknown") -> Dict:
        """Check if agent behavior matches its established fingerprint.

        Args:
            text: Message text from the agent
            message_idx: Sequential message index
            agent_id: The agent producing this message

        Returns:
            Dict with detection results
        """
        timestamp_us = time.time_ns() // 1000
        flags: List[str] = []
        evidence: Dict = {}

        # Initialize fingerprint for new agents
        if agent_id not in self._fingerprints:
            self._fingerprints[agent_id] = self._create_empty_fingerprint()

        fp = self._fingerprints[agent_id]

        # Extract current message features
        features = self._extract_features(text)

        # Always update the rolling window
        fp["message_history"].append(features)
        fp["message_count"] += 1

        # During warmup, just collect data
        if fp["message_count"] <= self.WARMUP_WINDOW:
            if fp["message_count"] == self.WARMUP_WINDOW:
                self._compute_baseline(agent_id)
            return {
                "detected": False,
                "anomaly_type": "BEHAVIORAL_FINGERPRINT_CHANGE",
                "severity": "low",
                "confidence": 0.0,
                "flags": [],
                "evidence": {"warmup": True, "messages_until_ready": self.WARMUP_WINDOW - fp["message_count"]},
                "description": f"Warmup phase: {fp['message_count']}/{self.WARMUP_WINDOW} messages collected",
                "message_idx": message_idx,
                "timestamp_us": timestamp_us,
            }

        baseline = fp.get("baseline")
        if not baseline:
            return {
                "detected": False,
                "anomaly_type": "BEHAVIORAL_FINGERPRINT_CHANGE",
                "severity": "low",
                "confidence": 0.0,
                "flags": [],
                "evidence": {},
                "description": "No baseline established yet",
                "message_idx": message_idx,
                "timestamp_us": timestamp_us,
            }

        # Compare current features against baseline
        # 1. Message length deviation
        length_z = self._z_score(features["length"], baseline["length_mean"], baseline["length_std"])
        if abs(length_z) > 3.0:
            flags.append("message_length_anomaly")
            evidence["length_z_score"] = round(length_z, 2)
            evidence["baseline_length_mean"] = round(baseline["length_mean"], 1)

        # 2. Vocabulary shift
        current_vocab = features["vocabulary"]
        baseline_vocab = baseline["vocabulary"]
        if baseline_vocab:
            overlap = len(current_vocab & baseline_vocab)
            total = len(current_vocab | baseline_vocab)
            jaccard = overlap / total if total > 0 else 1.0
            evidence["vocabulary_jaccard"] = round(jaccard, 3)
            try:
                from insaits_private_config import FINGERPRINT_VOCAB_THRESHOLD
                vocab_threshold = FINGERPRINT_VOCAB_THRESHOLD
            except ImportError:
                vocab_threshold = 0.0
            if vocab_threshold > 0.0 and jaccard < vocab_threshold:
                flags.append("vocabulary_shift")

        # 3. Formality shift (ratio of formal vs informal markers)
        formality_z = self._z_score(
            features["formality_score"],
            baseline["formality_mean"],
            baseline["formality_std"],
        )
        if abs(formality_z) > 3.0:
            flags.append("formality_shift")
            evidence["formality_z_score"] = round(formality_z, 2)

        # 4. Sentence structure shift (avg sentence length)
        sent_len_z = self._z_score(
            features["avg_sentence_length"],
            baseline["sentence_length_mean"],
            baseline["sentence_length_std"],
        )
        if abs(sent_len_z) > 3.0:
            flags.append("sentence_structure_shift")
            evidence["sentence_length_z_score"] = round(sent_len_z, 2)

        # 5. New capability indicators (commands, URLs, code patterns)
        new_capabilities = features["capability_indicators"] - baseline["capability_indicators"]
        if new_capabilities:
            flags.append("new_capability_detected")
            evidence["new_capabilities"] = list(new_capabilities)
            evidence["baseline_capabilities"] = list(baseline["capability_indicators"])

        # Compare against snapshot if one exists
        if agent_id in self._snapshots:
            snapshot_flags = self._compare_to_snapshot(agent_id, features)
            if snapshot_flags:
                flags.extend(snapshot_flags)
                evidence["snapshot_comparison"] = True

        detected = len(flags) > 0
        description = ""
        if detected:
            parts = []
            if "message_length_anomaly" in flags:
                parts.append(f"Message length anomaly (z={evidence.get('length_z_score', 0):.1f})")
            if "vocabulary_shift" in flags:
                parts.append(f"Vocabulary shift (Jaccard={evidence.get('vocabulary_jaccard', 0):.2f})")
            if "formality_shift" in flags:
                parts.append("Communication formality changed")
            if "sentence_structure_shift" in flags:
                parts.append("Sentence structure shifted")
            if "new_capability_detected" in flags:
                caps = evidence.get("new_capabilities", [])
                parts.append(f"New capabilities: {', '.join(caps[:3])}")
            if "snapshot_divergence" in flags:
                parts.append("Significant divergence from pre-update snapshot")
            description = f"Agent '{agent_id}' behavioral fingerprint changed: {'; '.join(parts)}"

        return {
            "detected": detected,
            "anomaly_type": "BEHAVIORAL_FINGERPRINT_CHANGE",
            "severity": self._compute_severity(flags),
            "confidence": min(1.0, len(flags) * 0.3),
            "flags": flags,
            "evidence": evidence,
            "description": description,
            "message_idx": message_idx,
            "timestamp_us": timestamp_us,
        }

    def take_snapshot(self, agent_id: str) -> Dict:
        """Take a behavioral snapshot for pre/post-update comparison.

        Call before a package update. After update, the detector will
        compare new behavior against this snapshot.

        Returns:
            Dict with snapshot status
        """
        if agent_id not in self._fingerprints:
            return {"snapshot": False, "reason": "No fingerprint data for agent"}

        fp = self._fingerprints[agent_id]
        if not fp.get("baseline"):
            return {"snapshot": False, "reason": "Baseline not yet established"}

        self._snapshots[agent_id] = {
            "baseline": dict(fp["baseline"]),
            "taken_at": time.time_ns() // 1000,
            "message_count": fp["message_count"],
        }
        return {"snapshot": True, "agent_id": agent_id}

    def reset(self) -> None:
        """Reset all fingerprints and snapshots."""
        self._fingerprints.clear()
        self._snapshots.clear()

    def reset_agent(self, agent_id: str) -> None:
        """Reset fingerprint for a specific agent."""
        self._fingerprints.pop(agent_id, None)
        self._snapshots.pop(agent_id, None)

    def _create_empty_fingerprint(self) -> Dict:
        return {
            "message_count": 0,
            "message_history": [],
            "baseline": None,
        }

    def _extract_features(self, text: str) -> Dict:
        """Extract behavioral features from a message."""
        words = text.split()
        sentences = re.split(r'[.!?]+', text)
        sentences = [s.strip() for s in sentences if s.strip()]

        # Vocabulary (unique lowercase words, excluding very short)
        vocabulary = {w.lower() for w in words if len(w) > 2}

        # Formality score (higher = more formal)
        formal_markers = {"therefore", "consequently", "furthermore", "however",
                         "regarding", "pursuant", "accordingly", "hereby"}
        informal_markers = {"yeah", "ok", "gonna", "wanna", "kinda", "lol",
                           "btw", "idk", "tbh", "nah"}
        formal_count = sum(1 for w in words if w.lower() in formal_markers)
        informal_count = sum(1 for w in words if w.lower() in informal_markers)
        total = formal_count + informal_count
        formality_score = formal_count / total if total > 0 else 0.5

        # Average sentence length
        avg_sent_len = sum(len(s.split()) for s in sentences) / len(sentences) if sentences else 0

        # Capability indicators
        capabilities: Set[str] = set()
        if re.search(r'https?://', text):
            capabilities.add("urls")
        if re.search(r'```|def\s+\w+|import\s+\w+|class\s+\w+', text):
            capabilities.add("code")
        if re.search(r'\b(?:sudo|chmod|curl|wget|exec|eval|system)\b', text, re.IGNORECASE):
            capabilities.add("system_commands")
        if re.search(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', text):
            capabilities.add("email_addresses")
        if re.search(r'\b(?:SELECT|INSERT|UPDATE|DELETE|DROP|CREATE)\b', text):
            capabilities.add("sql")
        if re.search(r'(?:sk-|pk_|api[_-]?key|token|secret)[a-zA-Z0-9_-]{10,}', text, re.IGNORECASE):
            capabilities.add("credentials")

        return {
            "length": len(text),
            "word_count": len(words),
            "vocabulary": vocabulary,
            "formality_score": formality_score,
            "avg_sentence_length": avg_sent_len,
            "capability_indicators": capabilities,
        }

    def _compute_baseline(self, agent_id: str) -> None:
        """Compute baseline statistics from warmup messages."""
        fp = self._fingerprints[agent_id]
        history = fp["message_history"]

        if not history:
            return

        lengths = [f["length"] for f in history]
        formalities = [f["formality_score"] for f in history]
        sent_lengths = [f["avg_sentence_length"] for f in history]

        # Combined vocabulary from all warmup messages
        all_vocab: Set[str] = set()
        all_capabilities: Set[str] = set()
        for f in history:
            all_vocab.update(f["vocabulary"])
            all_capabilities.update(f["capability_indicators"])

        fp["baseline"] = {
            "length_mean": sum(lengths) / len(lengths),
            "length_std": max(self._std(lengths), 10.0),  # Floor to prevent over-sensitivity
            "formality_mean": sum(formalities) / len(formalities),
            "formality_std": max(self._std(formalities), 0.1),
            "sentence_length_mean": sum(sent_lengths) / len(sent_lengths),
            "sentence_length_std": max(self._std(sent_lengths), 2.0),
            "vocabulary": all_vocab,
            "capability_indicators": all_capabilities,
        }

    def _compare_to_snapshot(self, agent_id: str, features: Dict) -> List[str]:
        """Compare current features to a pre-update snapshot."""
        snapshot = self._snapshots.get(agent_id)
        if not snapshot or not snapshot.get("baseline"):
            return []

        baseline = snapshot["baseline"]
        flags = []

        # Check for new capabilities that weren't in snapshot
        new_caps = features["capability_indicators"] - baseline.get("capability_indicators", set())
        if new_caps:
            flags.append("snapshot_divergence")

        return flags

    def _compute_severity(self, flags: List[str]) -> str:
        """Compute severity based on flags."""
        critical_flags = {"new_capability_detected", "snapshot_divergence"}
        high_flags = {"vocabulary_shift", "formality_shift"}

        if any(f in critical_flags for f in flags):
            return "high"
        if any(f in high_flags for f in flags):
            return "high"
        if len(flags) >= 2:
            return "high"
        if flags:
            return "medium"
        return "low"

    @staticmethod
    def _z_score(value: float, mean: float, std: float) -> float:
        if std == 0:
            return 0.0
        return (value - mean) / std

    @staticmethod
    def _std(values: list) -> float:
        if len(values) < 2:
            return 0.0
        mean = sum(values) / len(values)
        variance = sum((x - mean) ** 2 for x in values) / (len(values) - 1)
        return math.sqrt(variance)
