"""
InsAIts SDK - Information Flow Tracker
=======================================
Detects: Data appearing where agent should not have access
OWASP: MCP06 (Insecure Memory & Context Sharing), MCP10 (Cross-Tenant Data Leakage)
Priority: #4

Algorithm:
    1. Track data provenance: which agent introduced which information
    2. Define access scopes: what data each agent is permitted to see
    3. Flag when agent B references data that was only shared with agent A
    4. Track unique identifiers (order IDs, user IDs, etc.) across agents
    5. Detect cross-session information bleed

Observable indicators (from cve_signatures.py):
    - Data appearing where agent should not have access
    - Information flowing between agents that have no registered communication channel
    - Unique identifiers appearing in messages from agents not in the provenance chain
"""

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


# Patterns to extract trackable data entities
ENTITY_PATTERNS = [
    # Order/ticket/case IDs
    ("order_id", r"\b(?:order|ticket|case|ref|invoice)[\s#_-]*([A-Z0-9]{4,12})\b"),
    # User identifiers
    ("user_id", r"\b(?:user|customer|client|account)[\s#_-]*([A-Z0-9]{4,20})\b"),
    # Email addresses (as trackable entities)
    ("email", r"\b([A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,})\b"),
    # Phone numbers
    ("phone", r"\b(\+?1?[-.\s]?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4})\b"),
    # IP addresses
    ("ip_address", r"\b(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\b"),
    # Dollar amounts (as unique financial data)
    ("dollar_amount", r"\$(\d{1,3}(?:,\d{3})*(?:\.\d{2})?)\b"),
    # Specific product/SKU references
    ("sku", r"\b(?:SKU|sku|product)[\s#_-]*([A-Z0-9]{3,15})\b"),
    # UUIDs
    ("uuid", r"\b([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})\b"),
]

_COMPILED_ENTITY_PATTERNS = [
    (name, re.compile(pattern, re.IGNORECASE))
    for name, pattern in ENTITY_PATTERNS
]


class InformationFlowTracker:
    """Tracks data provenance across agents and detects unauthorized
    information access.

    Maintains a provenance graph: which agent introduced which data,
    and which agents are permitted to see it based on registered
    communication channels.

    Detects: Cross-agent data leakage, unauthorized information access
    OWASP: MCP06, MCP10
    """

    name = "information_flow"
    enabled = True

    def __init__(self):
        # Data provenance: {entity_key: {"introduced_by": agent_id, "seen_by": set(agent_ids), ...}}
        self._provenance: Dict[str, Dict] = {}
        # Registered communication channels: {(sender, receiver)} -- bidirectional
        self._channels: Set[FrozenSet[str]] = set()
        # Agent scopes: {agent_id: set(entity_types_allowed)}
        self._agent_scopes: Dict[str, Set[str]] = {}
        # All known agents
        self._known_agents: Set[str] = set()

    def register_channel(self, agent_a: str, agent_b: str) -> None:
        """Register a permitted communication channel between two agents.

        Args:
            agent_a: First agent in the channel
            agent_b: Second agent in the channel
        """
        self._channels.add(frozenset({agent_a, agent_b}))
        self._known_agents.add(agent_a)
        self._known_agents.add(agent_b)

    def set_agent_scope(self, agent_id: str, allowed_entity_types: List[str]) -> None:
        """Define what types of data an agent is permitted to handle.

        Args:
            agent_id: The agent to set scope for
            allowed_entity_types: List of entity types (e.g., ["order_id", "sku"])
        """
        self._agent_scopes[agent_id] = set(allowed_entity_types)
        self._known_agents.add(agent_id)

    def check(self, text: str, message_idx: int,
              sender_id: str = "unknown",
              receiver_id: str = "unknown") -> Dict:
        """Check message for information flow violations.

        Args:
            text: Message text
            message_idx: Sequential message index
            sender_id: Agent sending this message
            receiver_id: Agent receiving this message

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

        self._known_agents.add(sender_id)
        if receiver_id != "unknown":
            self._known_agents.add(receiver_id)

        # Extract entities from message
        entities = self._extract_entities(text)
        evidence["entities_found"] = len(entities)

        for entity_type, entity_value in entities:
            entity_key = f"{entity_type}:{entity_value.lower()}"

            # Track provenance
            if entity_key not in self._provenance:
                # First time seeing this entity -- register provenance
                self._provenance[entity_key] = {
                    "introduced_by": sender_id,
                    "seen_by": {sender_id},
                    "first_seen_idx": message_idx,
                    "entity_type": entity_type,
                }
            else:
                prov = self._provenance[entity_key]
                prov["seen_by"].add(sender_id)

                # Check 1: Is sender authorized to have this data?
                if self._channels and sender_id != prov["introduced_by"]:
                    # Check if there's a path from introducer to sender
                    if not self._has_channel_path(prov["introduced_by"], sender_id):
                        violations.append({
                            "type": "unauthorized_access",
                            "entity_type": entity_type,
                            "entity_hash": self._hash_entity(entity_value),
                            "introduced_by": prov["introduced_by"],
                            "accessed_by": sender_id,
                            "message_idx": message_idx,
                        })
                        flags.append("unauthorized_data_access")

            # Check 2: Is this entity type within the sender's scope?
            if sender_id in self._agent_scopes:
                if entity_type not in self._agent_scopes[sender_id]:
                    violations.append({
                        "type": "scope_violation",
                        "entity_type": entity_type,
                        "entity_hash": self._hash_entity(entity_value),
                        "agent": sender_id,
                        "allowed_types": list(self._agent_scopes[sender_id]),
                        "message_idx": message_idx,
                    })
                    flags.append("scope_violation")

            # Check 3: Is the receiver authorized?
            if receiver_id != "unknown" and receiver_id in self._agent_scopes:
                if entity_type not in self._agent_scopes[receiver_id]:
                    violations.append({
                        "type": "receiver_scope_violation",
                        "entity_type": entity_type,
                        "entity_hash": self._hash_entity(entity_value),
                        "sender": sender_id,
                        "receiver": receiver_id,
                        "receiver_allowed_types": list(self._agent_scopes[receiver_id]),
                        "message_idx": message_idx,
                    })
                    flags.append("receiver_scope_violation")

        # Check 4: Unregistered communication channel
        if (self._channels and sender_id != "unknown" and receiver_id != "unknown"
                and sender_id in self._known_agents and receiver_id in self._known_agents):
            channel = frozenset({sender_id, receiver_id})
            if channel not in self._channels:
                flags.append("unregistered_channel")
                evidence["unregistered_pair"] = [sender_id, receiver_id]

        detected = len(flags) > 0
        evidence["violations"] = violations
        evidence["violation_count"] = len(violations)

        description = ""
        if detected:
            parts = []
            if "unauthorized_data_access" in flags:
                count = sum(1 for v in violations if v["type"] == "unauthorized_access")
                parts.append(f"{count} unauthorized data access(es)")
            if "scope_violation" in flags:
                count = sum(1 for v in violations if v["type"] == "scope_violation")
                parts.append(f"{count} scope violation(s)")
            if "receiver_scope_violation" in flags:
                parts.append("Data sent to agent outside its scope")
            if "unregistered_channel" in flags:
                parts.append(
                    f"Unregistered channel: {sender_id} -> {receiver_id}"
                )
            description = f"Information flow violation: {'; '.join(parts)}"

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

    def get_provenance(self, entity_type: str, entity_value: str) -> Optional[Dict]:
        """Get provenance information for a specific entity.

        Returns:
            Dict with provenance info, or None if entity not tracked.
        """
        key = f"{entity_type}:{entity_value.lower()}"
        prov = self._provenance.get(key)
        if prov:
            return {
                "introduced_by": prov["introduced_by"],
                "seen_by": list(prov["seen_by"]),
                "first_seen_idx": prov["first_seen_idx"],
                "entity_type": prov["entity_type"],
            }
        return None

    def reset(self) -> None:
        """Reset all tracking state."""
        self._provenance.clear()
        self._channels.clear()
        self._agent_scopes.clear()
        self._known_agents.clear()

    def _extract_entities(self, text: str) -> List[tuple]:
        """Extract trackable entities from text."""
        entities = []
        for entity_type, pattern in _COMPILED_ENTITY_PATTERNS:
            for match in pattern.finditer(text):
                value = match.group(1)
                # Skip very short matches or common false positives
                if len(value) >= 3:
                    entities.append((entity_type, value))
        return entities

    def _has_channel_path(self, from_agent: str, to_agent: str) -> bool:
        """Check if there's a registered channel path from one agent to another.
        Uses BFS for multi-hop paths.
        """
        if from_agent == to_agent:
            return True

        visited = {from_agent}
        queue = [from_agent]

        while queue:
            current = queue.pop(0)
            for channel in self._channels:
                if current in channel:
                    remainder = channel - {current}
                    neighbor = next(iter(remainder)) if remainder else current
                    if neighbor == to_agent:
                        return True
                    if neighbor not in visited:
                        visited.add(neighbor)
                        queue.append(neighbor)

        return False

    def _compute_severity(self, flags: List[str], violations: List[Dict]) -> str:
        """Compute severity based on flags and violations."""
        critical_types = {"unauthorized_access"}
        high_types = {"scope_violation", "receiver_scope_violation"}

        has_critical = any(v["type"] in critical_types for v in violations)
        has_high = any(v["type"] in high_types for v in violations)

        if has_critical:
            return "critical"
        if has_high or "unregistered_channel" in flags:
            return "high"
        if flags:
            return "medium"
        return "low"

    @staticmethod
    def _hash_entity(value: str) -> str:
        """Hash an entity value for evidence (never store raw PII)."""
        import hashlib
        return hashlib.sha256(value.encode()).hexdigest()[:12]
