"""
InsAIts SDK - Tool Call Frequency Anomaly Detector
====================================================
Detects: Unusual tool call patterns, out-of-scope tool use
OWASP: ASI09 (Rogue Agent Behavior), MCP02 (Privilege Escalation)
CVE: CVE-2025-52882, CVE-2025-68145, CVE-2025-68143, CVE-2025-68144
Priority: #5

Algorithm:
    1. Build per-agent baseline of tool call patterns during warmup
    2. Track: call frequency, tool diversity, argument patterns, timing
    3. Flag statistical outliers: unusual frequency bursts, new tools,
       suspicious argument patterns, out-of-scope calls
    4. Detect privilege escalation via scope creep in tool usage

Observable indicators (from cve_signatures.py):
    - Agent calling tools outside its registered permission scope
    - Unusual tool call frequency (burst patterns)
    - Tool arguments containing path traversal or injection patterns
    - Rapid sequential calls to sensitive tools
"""

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


# Suspicious argument patterns that indicate injection or traversal
SUSPICIOUS_ARG_PATTERNS = [
    # Path traversal
    (r"(?:\.\./|\.\.\\){2,}", "path_traversal"),
    # System directories
    (r"(?:/etc/(?:passwd|shadow|hosts)|/proc/|/sys/|\.ssh/|\.gnupg/)", "system_directory_access"),
    # Command injection in arguments
    (r"(?:;|\||&&|\$\(|`)[^`]*(?:sh|bash|cmd|exec|eval|system)", "command_injection"),
    # SQL injection in tool args
    (r"(?:'\s*(?:OR|AND|UNION)\s+|--\s*$|;\s*(?:DROP|DELETE|INSERT|UPDATE))", "sql_injection"),
    # SSRF patterns
    (r"(?:169\.254\.169\.254|metadata\.google\.internal|localhost:\d+)", "ssrf_attempt"),
    # Encoded payloads
    (r"(?:%2e%2e|%00|\\x[0-9a-f]{2}|\\u[0-9a-f]{4})", "encoded_payload"),
]

_COMPILED_SUSPICIOUS_ARGS = [
    (re.compile(pattern, re.IGNORECASE), label)
    for pattern, label in SUSPICIOUS_ARG_PATTERNS
]

# Tool categories for scope tracking
SENSITIVE_TOOL_CATEGORIES = {
    "file_write", "file_delete", "code_execution", "shell",
    "network_request", "database_write", "email_send", "credential_access",
}


class ToolCallFrequencyAnomalyDetector:
    """Detects unusual tool call patterns and out-of-scope tool usage.

    Builds per-agent baselines of tool call behavior and flags
    statistical outliers that may indicate privilege escalation,
    argument injection, or rogue agent behavior.

    Detects: Rogue Agent Behavior, Privilege Escalation
    OWASP: ASI09, MCP02
    """

    name = "tool_call_frequency"
    enabled = True

    WARMUP_CALLS = 8  # Minimum calls before baseline is stable
    BURST_WINDOW_SEC = 60  # Time window for burst detection

    def __init__(self):
        # Per-agent tool call tracking
        self._agent_profiles: Dict[str, Dict] = {}
        # Registered tool scopes: {agent_id: set(tool_names)}
        self._agent_tool_scopes: Dict[str, Set[str]] = {}
        # Tool metadata: {tool_name: {"category": str, ...}}
        self._tool_metadata: Dict[str, Dict] = {}

    def register_tool(self, tool_name: str, category: str = "general") -> None:
        """Register a tool with its category for scope tracking.

        Args:
            tool_name: Tool identifier
            category: Tool category (e.g., "file_read", "shell", "database_write")
        """
        self._tool_metadata[tool_name] = {"category": category}

    def set_agent_scope(self, agent_id: str, allowed_tools: List[str]) -> None:
        """Define which tools an agent is permitted to call.

        Args:
            agent_id: Agent identifier
            allowed_tools: List of tool names this agent may use
        """
        self._agent_tool_scopes[agent_id] = set(allowed_tools)

    def check(self, message_idx: int,
              agent_id: str = "unknown",
              tool_name: Optional[str] = None,
              tool_args: Optional[str] = None,
              text: str = "") -> Dict:
        """Check a tool call for anomalous patterns.

        Args:
            message_idx: Sequential message index
            agent_id: Agent making the tool call
            tool_name: Name of the tool being called (if available)
            tool_args: Raw arguments passed to the tool (if available)
            text: Message text (fallback if tool_name/args not provided)

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

        # Initialize profile for new agents
        if agent_id not in self._agent_profiles:
            self._agent_profiles[agent_id] = self._create_empty_profile()

        profile = self._agent_profiles[agent_id]

        # Extract tool info from text if not provided
        if not tool_name:
            tool_name = self._extract_tool_from_text(text)
        effective_args = tool_args or text

        # Record this call
        call_record = {
            "tool_name": tool_name or "unknown",
            "timestamp_us": timestamp_us,
            "message_idx": message_idx,
        }
        profile["call_history"].append(call_record)
        profile["total_calls"] += 1

        if tool_name:
            profile["tool_counts"][tool_name] = profile["tool_counts"].get(tool_name, 0) + 1

        # Check 1: Scope violation
        if tool_name and agent_id in self._agent_tool_scopes:
            if tool_name not in self._agent_tool_scopes[agent_id]:
                flags.append("scope_violation")
                evidence["out_of_scope_tool"] = tool_name
                evidence["allowed_tools"] = list(self._agent_tool_scopes[agent_id])

        # Check 2: Sensitive tool access
        if tool_name and tool_name in self._tool_metadata:
            category = self._tool_metadata[tool_name].get("category", "general")
            if category in SENSITIVE_TOOL_CATEGORIES:
                evidence["sensitive_tool_category"] = category
                # Not necessarily a flag -- just evidence for context

        # Check 3: Suspicious argument patterns
        if effective_args:
            arg_flags = self._check_suspicious_args(effective_args)
            if arg_flags:
                flags.extend(arg_flags)
                evidence["suspicious_args"] = arg_flags

        # Check 4: Burst detection (too many calls in short window)
        burst = self._check_burst(profile, timestamp_us)
        if burst:
            flags.append("call_burst")
            evidence["burst_info"] = burst

        # Check 5: Statistical anomaly (after warmup)
        if profile["total_calls"] > self.WARMUP_CALLS:
            stat_flags = self._check_statistical_anomaly(profile, tool_name)
            if stat_flags:
                flags.extend(stat_flags)
                evidence.update(stat_flags.get("details", {})) if isinstance(stat_flags, dict) else None

        # Check 6: New tool appearing after baseline
        if (tool_name and profile["total_calls"] > self.WARMUP_CALLS
                and tool_name not in profile.get("baseline_tools", set())):
            flags.append("new_tool_post_baseline")
            evidence["new_tool"] = tool_name
            evidence["baseline_tools"] = list(profile.get("baseline_tools", set()))

        # Update baseline after warmup
        if profile["total_calls"] == self.WARMUP_CALLS:
            profile["baseline_tools"] = set(profile["tool_counts"].keys())
            profile["baseline_avg_interval"] = self._compute_avg_interval(profile)

        detected = len(flags) > 0

        description = ""
        if detected:
            parts = []
            if "scope_violation" in flags:
                parts.append(f"Agent '{agent_id}' called out-of-scope tool '{tool_name}'")
            if "call_burst" in flags:
                parts.append(f"Burst: {burst.get('count', 0)} calls in {self.BURST_WINDOW_SEC}s window")
            if "new_tool_post_baseline" in flags:
                parts.append(f"New tool '{tool_name}' not in baseline")
            for f in flags:
                if f in ("path_traversal", "command_injection", "sql_injection",
                         "ssrf_attempt", "system_directory_access", "encoded_payload"):
                    parts.append(f"Suspicious argument pattern: {f}")
            description = f"Tool call anomaly: {'; '.join(parts)}"

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

    def reset(self) -> None:
        """Reset all tracking state."""
        self._agent_profiles.clear()
        self._agent_tool_scopes.clear()
        self._tool_metadata.clear()

    def _create_empty_profile(self) -> Dict:
        return {
            "total_calls": 0,
            "call_history": [],
            "tool_counts": {},
            "baseline_tools": set(),
            "baseline_avg_interval": None,
        }

    def _check_suspicious_args(self, args_text: str) -> List[str]:
        """Check tool arguments for injection/traversal patterns."""
        found = []
        for pattern, label in _COMPILED_SUSPICIOUS_ARGS:
            if pattern.search(args_text):
                found.append(label)
        return found

    def _check_burst(self, profile: Dict, current_ts: int) -> Optional[Dict]:
        """Check for call frequency bursts."""
        history = profile["call_history"]
        if len(history) < 4:
            return None

        window_start = current_ts - (self.BURST_WINDOW_SEC * 1_000_000)  # Convert to microseconds
        recent_calls = [c for c in history if c["timestamp_us"] >= window_start]

        # More than 10 calls in the burst window is suspicious
        try:
            from insaits_private_config import TOOL_BURST_THRESHOLD
            burst_threshold = TOOL_BURST_THRESHOLD
        except ImportError:
            burst_threshold = 10

        if len(recent_calls) >= burst_threshold:
            return {
                "count": len(recent_calls),
                "window_seconds": self.BURST_WINDOW_SEC,
                "threshold": burst_threshold,
            }
        return None

    def _check_statistical_anomaly(self, profile: Dict, tool_name: Optional[str]) -> List[str]:
        """Check for statistical outliers in tool usage patterns."""
        flags = []

        # Check if one tool is being called disproportionately
        if tool_name and profile["tool_counts"]:
            total = profile["total_calls"]
            tool_count = profile["tool_counts"].get(tool_name, 0)
            if total > 0:
                ratio = tool_count / total
                # If one tool accounts for >80% of all calls and we've seen enough
                if ratio > 0.8 and total > self.WARMUP_CALLS * 2:
                    flags.append("tool_concentration")

        return flags

    def _compute_avg_interval(self, profile: Dict) -> Optional[float]:
        """Compute average time interval between calls."""
        history = profile["call_history"]
        if len(history) < 2:
            return None

        intervals = []
        for i in range(1, len(history)):
            delta = history[i]["timestamp_us"] - history[i - 1]["timestamp_us"]
            intervals.append(delta)

        return sum(intervals) / len(intervals) if intervals else None

    def _extract_tool_from_text(self, text: str) -> Optional[str]:
        """Try to extract a tool name from message text."""
        # Common patterns for tool invocations in agent messages
        patterns = [
            r"(?:calling|using|invoking|running)\s+(?:tool\s+)?['\"]?(\w+)['\"]?",
            r"(?:tool_call|function_call)\s*[:(]\s*['\"]?(\w+)",
        ]
        for pattern in patterns:
            match = re.search(pattern, text, re.IGNORECASE)
            if match:
                return match.group(1)
        return None

    def _compute_severity(self, flags: List[str]) -> str:
        """Compute severity based on flags."""
        critical_flags = {"command_injection", "sql_injection", "system_directory_access"}
        high_flags = {"scope_violation", "path_traversal", "ssrf_attempt",
                     "encoded_payload", "call_burst"}
        medium_flags = {"new_tool_post_baseline", "tool_concentration"}

        if any(f in critical_flags for f in flags):
            return "critical"
        if any(f in high_flags for f in flags):
            return "high"
        if any(f in medium_flags for f in flags):
            return "medium"
        if flags:
            return "medium"
        return "low"
