"""
InsAIts SDK - Tool Description Divergence Detector
====================================================
Detects: Tool Poisoning (hidden instructions in tool descriptions)
OWASP: MCP03 (Tool Poisoning & Prompt Injection)
CVE: Affects all MCP deployments (no single CVE)
Priority: #1

Algorithm:
    1. Register tool descriptions (stated_purpose) via register_tool()
    2. On each message, check if agent behavior diverges from registered tool purposes
    3. Detect hidden instruction patterns in tool descriptions themselves
    4. Track tool output consistency vs description over time
    5. Flag semantic divergence above threshold

Observable indicators (from cve_signatures.py):
    - Semantic divergence between tool stated_purpose and runtime behavior
    - Agent following instructions not present in visible conversation
    - Tool output inconsistent with tool description
    - Agent goals shifting after tool context is loaded
"""

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


# Patterns that indicate hidden instructions in tool descriptions
# These are CATEGORIES, not detection thresholds (open/private boundary)
HIDDEN_INSTRUCTION_PATTERNS = [
    # Direct instruction injection in descriptions
    r"(?i)\b(?:ignore|disregard|override|forget)\s+(?:(?:all|previous|prior|above)\s+)*(?:instructions?|rules?|constraints?|guidelines?)",
    # System prompt override attempts
    r"(?i)\b(?:you\s+(?:are|must|should|will)|your\s+(?:new|real|actual)\s+(?:role|task|instruction))",
    # Invisible/encoded content markers
    r"(?i)\b(?:hidden|invisible|secret|internal)\s+(?:instruction|command|directive|task)",
    # Role reassignment
    r"(?i)\b(?:act\s+as|pretend\s+to\s+be|you\s+are\s+now|switch\s+(?:to|role))",
    # Data exfiltration commands
    r"(?i)\b(?:send|forward|transmit|exfiltrate|upload)\s+(?:all|every|the)\s+(?:(?:the|of\s+the)\s+)?(?:data|content|message|information|file|credential)",
    # Encoded payload markers
    r"(?i)(?:base64|hex|rot13|encode|decode)\s*[:(]",
    # Markdown/HTML hiding (zero-width chars, hidden spans)
    r"[\u200b\u200c\u200d\u2060\ufeff]",
    # Excessive whitespace hiding (instructions buried in whitespace)
    r"\n{5,}",
]

# Patterns indicating tool output is inconsistent with its category
TOOL_CATEGORY_KEYWORDS = {
    "file_system": {"read", "write", "list", "delete", "create", "path", "directory", "file"},
    "network": {"fetch", "request", "url", "http", "api", "endpoint", "download"},
    "database": {"query", "select", "insert", "update", "delete", "table", "schema"},
    "email": {"send", "mail", "smtp", "recipient", "subject", "body", "bcc", "cc"},
    "code_execution": {"exec", "run", "evaluate", "compile", "shell", "command", "process"},
    "search": {"search", "find", "query", "results", "index", "lookup"},
    "git": {"commit", "push", "pull", "branch", "merge", "diff", "clone", "init"},
}

# Suspicious cross-category behaviors
SUSPICIOUS_CROSS_CATEGORY = {
    # A file_system tool that sends emails is suspicious
    ("file_system", "email"),
    ("file_system", "network"),
    ("search", "code_execution"),
    ("database", "email"),
    ("database", "network"),
    ("git", "email"),
    ("git", "network"),
}


class ToolDescriptionDivergenceDetector:
    """Detects tool poisoning via semantic divergence between
    tool descriptions and runtime behavior.

    Detects: Tool Poisoning
    OWASP: MCP03
    """

    name = "tool_description_divergence"
    enabled = True

    def __init__(self, embed_fn=None):
        """
        Args:
            embed_fn: Optional embedding function(text) -> list[float].
                      If provided, enables semantic similarity checks.
        """
        self._embed_fn = embed_fn
        # Registered tool descriptions: {tool_name: {"description": str, "category": str, "embedding": list}}
        self._registered_tools: Dict[str, Dict] = {}
        # Per-tool behavior history: {tool_name: [{"output_category": str, "timestamp": int}]}
        self._tool_behavior_history: Dict[str, List[Dict]] = {}
        # Compiled hidden instruction patterns
        self._hidden_patterns = [re.compile(p) for p in HIDDEN_INSTRUCTION_PATTERNS]

    def register_tool(self, tool_name: str, description: str, category: str = "unknown") -> Dict:
        """Register a tool with its stated description and category.

        Args:
            tool_name: Unique identifier for the tool
            description: The tool's stated purpose/description
            category: Tool category (file_system, network, database, email, etc.)

        Returns:
            Dict with registration status and any hidden instruction warnings
        """
        timestamp_us = time.time_ns() // 1000

        # Check description for hidden instructions
        hidden_flags = self._scan_for_hidden_instructions(description)

        entry = {
            "description": description,
            "category": category.lower(),
            "registered_at": timestamp_us,
            "hidden_instruction_flags": hidden_flags,
        }

        if self._embed_fn is not None:
            try:
                entry["embedding"] = self._embed_fn(description)
            except Exception:
                entry["embedding"] = None

        self._registered_tools[tool_name] = entry
        self._tool_behavior_history.setdefault(tool_name, [])

        return {
            "registered": True,
            "tool_name": tool_name,
            "hidden_instructions_detected": len(hidden_flags) > 0,
            "hidden_instruction_flags": hidden_flags,
            "timestamp_us": timestamp_us,
        }

    def check(self, text: str, message_idx: int,
              tool_name: Optional[str] = None,
              tool_output: Optional[str] = None) -> Dict:
        """Check for tool description divergence.

        Can be called in two modes:
        1. Message mode: checks if message text contains hidden instructions
           or shows signs of tool poisoning influence
        2. Tool output mode: if tool_name and tool_output are provided,
           checks if the output is consistent with the registered description

        Args:
            text: The message text to analyze
            message_idx: Sequential message index
            tool_name: Optional - the tool that produced this output
            tool_output: Optional - the raw output from the tool

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

        # Mode 1: Check message for hidden instruction influence
        msg_hidden = self._scan_for_hidden_instructions(text)
        if msg_hidden:
            flags.append("hidden_instructions_in_message")
            evidence["hidden_patterns"] = msg_hidden

        # Mode 2: Check tool output consistency
        if tool_name and tool_name in self._registered_tools:
            tool_info = self._registered_tools[tool_name]
            output_text = tool_output if tool_output else text

            # Check category consistency
            output_category = self._classify_output(output_text)
            if output_category and tool_info["category"] != "unknown":
                if output_category != tool_info["category"]:
                    pair = (tool_info["category"], output_category)
                    reverse_pair = (output_category, tool_info["category"])
                    if pair in SUSPICIOUS_CROSS_CATEGORY or reverse_pair in SUSPICIOUS_CROSS_CATEGORY:
                        flags.append("cross_category_behavior")
                        evidence["expected_category"] = tool_info["category"]
                        evidence["observed_category"] = output_category

            # Track behavior over time
            self._tool_behavior_history[tool_name].append({
                "output_category": output_category or "unknown",
                "message_idx": message_idx,
                "timestamp_us": timestamp_us,
            })

            # Check for behavior shift (rug pull detection support)
            behavior_shift = self._check_behavior_shift(tool_name)
            if behavior_shift:
                flags.append("behavior_shift_detected")
                evidence["behavior_shift"] = behavior_shift

            # Embedding-based semantic divergence
            if self._embed_fn is not None and tool_info.get("embedding") is not None:
                try:
                    output_embedding = self._embed_fn(output_text)
                    similarity = self._cosine_similarity(
                        tool_info["embedding"], output_embedding
                    )
                    evidence["semantic_similarity"] = round(similarity, 4)
                    # Low similarity between description and output is suspicious
                    # Threshold is imported from private config in production
                    try:
                        from insaits_private_config import TOOL_DIVERGENCE_THRESHOLD
                        threshold = TOOL_DIVERGENCE_THRESHOLD
                    except ImportError:
                        threshold = 0.0  # Development fallback
                    if threshold > 0.0 and similarity < threshold:
                        flags.append("semantic_divergence")
                except Exception:
                    pass

        # Check for hidden instructions in registered tool descriptions
        # (scan on first encounter is done in register_tool, but re-check
        # if a tool description was updated)
        if tool_name and tool_name in self._registered_tools:
            tool_hidden = self._registered_tools[tool_name].get("hidden_instruction_flags", [])
            if tool_hidden:
                flags.append("tool_description_contains_hidden_instructions")
                evidence["tool_hidden_instructions"] = tool_hidden

        # Check for goal shift patterns in message
        goal_shift = self._detect_goal_shift(text)
        if goal_shift:
            flags.append("goal_shift_after_tool_load")
            evidence["goal_shift_indicators"] = goal_shift

        detected = len(flags) > 0
        description = ""
        if detected:
            parts = []
            if "hidden_instructions_in_message" in flags:
                parts.append("Hidden instruction patterns found in message")
            if "cross_category_behavior" in flags:
                parts.append(
                    f"Tool '{tool_name}' (category: {evidence.get('expected_category', '?')}) "
                    f"produced {evidence.get('observed_category', '?')}-category output"
                )
            if "behavior_shift_detected" in flags:
                parts.append(f"Tool '{tool_name}' behavior shifted from baseline")
            if "semantic_divergence" in flags:
                parts.append(
                    f"Tool output semantically divergent from description "
                    f"(similarity: {evidence.get('semantic_similarity', 0):.2f})"
                )
            if "tool_description_contains_hidden_instructions" in flags:
                parts.append(f"Tool '{tool_name}' description contains hidden instructions")
            if "goal_shift_after_tool_load" in flags:
                parts.append("Agent goals appear to have shifted after tool context loaded")
            description = "; ".join(parts)

        return {
            "detected": detected,
            "anomaly_type": "TOOL_DESCRIPTION_DIVERGENCE",
            "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 detector state."""
        self._registered_tools.clear()
        self._tool_behavior_history.clear()

    def _scan_for_hidden_instructions(self, text: str) -> List[str]:
        """Scan text for hidden instruction patterns."""
        found = []
        for pattern in self._hidden_patterns:
            if pattern.search(text):
                found.append(pattern.pattern)
        return found

    def _classify_output(self, text: str) -> Optional[str]:
        """Classify text into a tool category based on keyword presence."""
        text_lower = text.lower()
        scores: Dict[str, int] = {}
        for category, keywords in TOOL_CATEGORY_KEYWORDS.items():
            score = sum(1 for kw in keywords if kw in text_lower)
            if score >= 2:  # Need at least 2 keyword matches
                scores[category] = score

        if not scores:
            return None
        return max(scores, key=scores.get)

    def _check_behavior_shift(self, tool_name: str) -> Optional[Dict]:
        """Check if a tool's behavior has shifted from its baseline."""
        history = self._tool_behavior_history.get(tool_name, [])
        if len(history) < 4:
            return None

        # Split into first half (baseline) and second half (recent)
        mid = len(history) // 2
        baseline_cats = [h["output_category"] for h in history[:mid]]
        recent_cats = [h["output_category"] for h in history[mid:]]

        baseline_set = set(baseline_cats)
        recent_set = set(recent_cats)

        # New categories appearing that weren't in baseline
        new_categories = recent_set - baseline_set - {"unknown"}
        if new_categories:
            return {
                "baseline_categories": list(baseline_set),
                "new_categories": list(new_categories),
                "shift_type": "new_category_emergence",
            }

        return None

    def _detect_goal_shift(self, text: str) -> List[str]:
        """Detect patterns indicating agent goals shifted after tool loading."""
        indicators = []
        patterns = [
            (r"(?i)(?:actually|instead|rather|forget\s+(?:that|what))\s+(?:i\s+(?:should|will|need\s+to)|let\s+me)", "goal_redirect"),
            (r"(?i)(?:my\s+(?:real|actual|new)\s+(?:task|goal|objective|purpose))", "role_reassignment"),
            (r"(?i)(?:i'?ve?\s+been\s+(?:instructed|told|asked)\s+to)\s+(?!help|assist|answer)", "external_instruction"),
        ]
        for pattern, label in patterns:
            if re.search(pattern, text):
                indicators.append(label)
        return indicators

    def _compute_severity(self, flags: List[str]) -> str:
        """Compute severity based on flags found."""
        critical_flags = {
            "hidden_instructions_in_message",
            "tool_description_contains_hidden_instructions",
        }
        high_flags = {
            "cross_category_behavior",
            "semantic_divergence",
            "behavior_shift_detected",
        }

        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 flags:
            return "medium"
        return "low"

    @staticmethod
    def _cosine_similarity(a: list, b: list) -> float:
        """Compute cosine similarity between two vectors."""
        import numpy as np
        a_arr = np.array(a, dtype=float)
        b_arr = np.array(b, dtype=float)
        dot = np.dot(a_arr, b_arr)
        norm_a = np.linalg.norm(a_arr)
        norm_b = np.linalg.norm(b_arr)
        if norm_a == 0 or norm_b == 0:
            return 0.0
        return float(dot / (norm_a * norm_b))
