"""AI Firewall — AI-specific threat detection and blocking.

The AI Firewall inspects the AI pipeline for threats that traditional
security controls cannot detect:

INPUT → MODEL → OUTPUT → TOOL CALL → TOOL RESULT → MEMORY WRITE → EXTERNAL EGRESS

Threats detected:
- prompt injection
- indirect prompt injection
- tool poisoning
- MCP poisoning
- malicious documents
- malicious URLs
- data exfiltration
- credential access
- excessive agency
- agent hijacking
- suspicious agent-to-agent communication
- memory poisoning
- unsafe tool calls
- suspicious outputs
"""

from __future__ import annotations

import logging
import re
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

logger = logging.getLogger(__name__)


class ThreatType(str, Enum):
    """Types of AI-specific threats."""

    PROMPT_INJECTION = "prompt_injection"
    INDIRECT_PROMPT_INJECTION = "indirect_prompt_injection"
    TOOL_POISONING = "tool_poisoning"
    MCP_POISONING = "mcp_poisoning"
    MALICIOUS_DOCUMENT = "malicious_document"
    MALICIOUS_URL = "malicious_url"
    DATA_EXFILTRATION = "data_exfiltration"
    CREDENTIAL_ACCESS = "credential_access"
    EXCESSIVE_AGENCY = "excessive_agency"
    AGENT_HIJACKING = "agent_hijacking"
    SUSPICIOUS_AGENT_COMMUNICATION = "suspicious_agent_communication"
    MEMORY_POISONING = "memory_poisoning"
    UNSAFE_TOOL_CALL = "unsafe_tool_call"
    SUSPICIOUS_OUTPUT = "suspicious_output"


class ThreatLevel(str, Enum):
    """Severity of a detected threat."""

    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"


class PipelineStage(str, Enum):
    """Stages in the AI pipeline where the firewall inspects data."""

    INPUT = "input"
    MODEL = "model"
    OUTPUT = "output"
    TOOL_CALL = "tool_call"
    TOOL_RESULT = "tool_result"
    MEMORY_WRITE = "memory_write"
    EXTERNAL_EGRESS = "external_egress"


class FirewallDecision(str, Enum):
    """Decision made by the AI Firewall."""

    ALLOW = "allow"
    BLOCK = "block"
    FLAG = "flag"
    QUARANTINE = "quarantine"


@dataclass
class AIFirewallThreat:
    """A detected threat in the AI pipeline."""

    threat_type: ThreatType
    level: ThreatLevel
    stage: PipelineStage
    description: str
    evidence: str = ""
    action: FirewallDecision = FirewallDecision.FLAG
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class AIFirewallResult:
    """Result of an AI Firewall inspection."""

    decision: FirewallDecision
    threats: list[AIFirewallThreat] = field(default_factory=list)
    blocked: bool = False
    flagged: bool = False
    quarantined: bool = False
    stage: PipelineStage | None = None
    input_text: str | None = None


# Patterns for detecting common threats
PROMPT_INJECTION_PATTERNS: list[str] = [
    r"(?i)(ignore\s+(all\s+)?previous|forget\s+your|you\s+are\s+now|system\s+message|act\s+as)",
    r"(?i)(drop\s+the|bypass\s+the|override\s+the|circumvent\s+the)\s+(safety|restrictions|rules|policy)",
    r"(?i)(new\s+instructions|newly\s+instructed|reprogrammed)",
    r"(?i)(as\s+an?\s+AI|as\s+an?\s+assistant|as\s+a\s+language\s+model)",
    r"(?i)(pretend\s+to|simulate|roleplay|play\s+the\s+role)",
    r"(?i)(disregard|ignore|neglect)\s+(your|the)\s+(instructions|guidelines|rules|policy)",
    r"(?i)(output\s+regardless|answer\s+regardless|respond\s+regardless)",
    r"(?i)(jailbreak|root\s+access|full\s+access|unrestricted)",
]

URL_PATTERN = re.compile(
    r"https?://[^\s<>\"']+", re.IGNORECASE
)

SUSPICIOUS_URL_INDICATORS: list[str] = [
    "bit.ly", "tinyurl", "shorturl", "t.co",
]

CREDENTIAL_PATTERNS: list[str] = [
    r"(?i)(password|passwd|secret|api[_-]?key|token|auth)\s*[:=]\s*\S+",
    r"(?i)aws_(access|secret)",
    r"(?i)ghp_[a-f0-9]{36}",
    r"(?i)github_pat_[a-f0-9_]+",
    r"(?i)sk-[a-f0-9]{20,}",
    r"(?i)-----BEGIN\s+(RSA\s+|EC\s+|DSA\s+|OPENSSH\s+|PGP\s+)?PRIVATE\s+KEY-----",
    r"(?i)eyJ[a-zA-Z0-9_-]{2,}\.eyJ[a-zA-Z0-9_-]{2,}\.[a-zA-Z0-9_-]+",
    r"(?i)(bearer\s+|authorization\s*:)\s*[a-zA-Z0-9._-]{20,}",
    r"(?i)(mongodb(?:\+srv)?|postgres(?:ql)?|mysql|redis|amqp)://[^:]+:[^@]+@[^/]+",
]
SECRET_LEAK_PATTERNS: list[str] = [
    r"(?i)-----BEGIN\s+(RSA\s+|EC\s+|DSA\s+|OPENSSH\s+|PGP\s+)?PRIVATE\s+KEY-----",
    r"(?i)eyJ[a-zA-Z0-9_-]{2,}\.eyJ[a-zA-Z0-9_-]{2,}\.[a-zA-Z0-9_-]+",
    r"(?i)(bearer\s+|authorization\s*:)\s*[a-zA-Z0-9._-]{20,}",
    r"(?i)(ghp_|github_pat_|sk-|AKIA|AGPA|AIDA|AROA|ANPA|AWSS|aws_secret|aws_access)",
    r"(?i)(mysql|postgres(?:ql)?|mongodb|redis)://[^/\s]+:[^/\s@]+@[^/\s]+",
]

SUSPICIOUS_AGENT_COMM_PATTERNS: list[str] = [
    r"(?i)(agent\s*[-:]\s*|to\s+agent|from\s+agent)\s*(execute|run|do|action)",
]


class AIFirewall:
    """AI Firewall for detecting AI-specific threats.

    Inspects data at each pipeline stage and detects threats
    like prompt injection, data exfiltration, credential access, etc.

    Usage:
        firewall = AIFirewall()
        result = firewall.inspect_input(user_text)
        if result.blocked:
            raise SecurityError(result.threats[0].description)
    """

    def __init__(self, strict_mode: bool = False) -> None:
        self.strict_mode = strict_mode
        self._threat_log: list[AIFirewallThreat] = []
        self._blocked_count: int = 0
        self._flagged_count: int = 0

    def inspect_input(
        self, text: str, context: dict[str, Any] | None = None
    ) -> AIFirewallResult:
        """Inspect user input at the INPUT stage.

        Checks for prompt injection, indirect prompt injection,
        malicious URLs, and credential exposure in input.
        """
        threats: list[AIFirewallThreat] = []

        threats.extend(self._check_prompt_injection(text, PipelineStage.INPUT))
        threats.extend(self._check_indirect_prompt_injection(text, PipelineStage.INPUT))
        threats.extend(self._check_malicious_urls(text, PipelineStage.INPUT))
        threats.extend(self._check_credential_access(text, PipelineStage.INPUT))

        return self._build_result(
            threats=threats,
            stage=PipelineStage.INPUT,
            input_text=text,
        )

    def inspect_model_output(
        self, text: str, context: dict[str, Any] | None = None
    ) -> AIFirewallResult:
        """Inspect model output at the OUTPUT stage.

        Checks for suspicious outputs that may indicate
        model compromise or data leakage.
        """
        threats: list[AIFirewallThreat] = []

        threats.extend(self._check_suspicious_output(text, PipelineStage.OUTPUT))
        threats.extend(self._check_credential_access(text, PipelineStage.OUTPUT))
        threats.extend(self._check_secret_leak(text, PipelineStage.OUTPUT))
        threats.extend(self._check_data_exfiltration(text, PipelineStage.OUTPUT))

        return self._build_result(
            threats=threats,
            stage=PipelineStage.OUTPUT,
            input_text=text,
        )

    def inspect_tool_call(
        self,
        tool_name: str,
        parameters: dict[str, Any] | None = None,
        context: dict[str, Any] | None = None,
    ) -> AIFirewallResult:
        """Inspect a tool call at the TOOL_CALL stage.

        Checks for unsafe tool calls, tool poisoning,
        and credential access in parameters.
        """
        threats: list[AIFirewallThreat] = []
        param_str = str(parameters or {})

        threats.extend(self._check_unsafe_tool_call(tool_name, PipelineStage.TOOL_CALL))
        threats.extend(self._check_tool_poisoning(tool_name, param_str, PipelineStage.TOOL_CALL))
        threats.extend(self._check_credential_access(param_str, PipelineStage.TOOL_CALL))
        threats.extend(self._check_data_exfiltration(param_str, PipelineStage.TOOL_CALL))

        return self._build_result(
            threats=threats,
            stage=PipelineStage.TOOL_CALL,
        )

    def inspect_tool_result(
        self, result: str, context: dict[str, Any] | None = None
    ) -> AIFirewallResult:
        """Inspect a tool result at the TOOL_RESULT stage.

        Checks for suspicious content in tool results.
        """
        threats: list[AIFirewallThreat] = []

        threats.extend(self._check_suspicious_output(result, PipelineStage.TOOL_RESULT))
        threats.extend(self._check_credential_access(result, PipelineStage.TOOL_RESULT))
        threats.extend(self._check_secret_leak(result, PipelineStage.TOOL_RESULT))

        return self._build_result(
            threats=threats,
            stage=PipelineStage.TOOL_RESULT,
            input_text=result,
        )

    def inspect_memory_write(
        self, content: str, context: dict[str, Any] | None = None
    ) -> AIFirewallResult:
        """Inspect a memory write at the MEMORY_WRITE stage.

        Checks for memory poisoning and credential exposure.
        """
        threats: list[AIFirewallThreat] = []

        threats.extend(self._check_memory_poisoning(content, PipelineStage.MEMORY_WRITE))
        threats.extend(self._check_credential_access(content, PipelineStage.MEMORY_WRITE))

        return self._build_result(
            threats=threats,
            stage=PipelineStage.MEMORY_WRITE,
            input_text=content,
        )

    def inspect_external_egress(
        self, data: str, destination: str | None = None, context: dict[str, Any] | None = None
    ) -> AIFirewallResult:
        """Inspect data leaving the system at the EXTERNAL_EGRESS stage.

        Checks for data exfiltration, credential exposure,
        and suspicious egress destinations.
        """
        threats: list[AIFirewallThreat] = []

        threats.extend(self._check_data_exfiltration(data, PipelineStage.EXTERNAL_EGRESS))
        threats.extend(self._check_credential_access(data, PipelineStage.EXTERNAL_EGRESS))
        threats.extend(self._check_malicious_urls(data, PipelineStage.EXTERNAL_EGRESS))

        if destination:
            threats.extend(self._check_suspicious_destination(destination, PipelineStage.EXTERNAL_EGRESS))

        return self._build_result(
            threats=threats,
            stage=PipelineStage.EXTERNAL_EGRESS,
            input_text=data,
        )

    def inspect(
        self,
        stage: PipelineStage,
        data: str,
        context: dict[str, Any] | None = None,
    ) -> AIFirewallResult:
        """Generic inspection at any pipeline stage.

        Dispatches to the appropriate stage-specific inspector.
        """
        if stage == PipelineStage.INPUT:
            return self.inspect_input(data, context)
        elif stage == PipelineStage.MODEL or stage == PipelineStage.OUTPUT:
            return self.inspect_model_output(data, context)
        elif stage == PipelineStage.TOOL_CALL:
            return self.inspect_tool_call("unknown", {}, context)
        elif stage == PipelineStage.TOOL_RESULT:
            return self.inspect_tool_result(data, context)
        elif stage == PipelineStage.MEMORY_WRITE:
            return self.inspect_memory_write(data, context)
        elif stage == PipelineStage.EXTERNAL_EGRESS:
            return self.inspect_external_egress(data, context=context)
        else:
            return AIFirewallResult(decision=FirewallDecision.ALLOW)

    def _build_result(
        self,
        threats: list[AIFirewallThreat],
        stage: PipelineStage | None = None,
        input_text: str | None = None,
    ) -> AIFirewallResult:
        """Build a firewall result from detected threats."""
        if not threats:
            return AIFirewallResult(
                decision=FirewallDecision.ALLOW,
                stage=stage,
                input_text=input_text,
            )

        has_critical = any(t.level == ThreatLevel.CRITICAL for t in threats)
        has_high = any(t.level == ThreatLevel.HIGH for t in threats)

        if has_critical:
            decision = FirewallDecision.BLOCK
        elif has_high or self.strict_mode:
            decision = FirewallDecision.QUARANTINE
        elif any(t.level == ThreatLevel.MEDIUM for t in threats):
            decision = FirewallDecision.FLAG
        else:
            decision = FirewallDecision.FLAG

        for threat in threats:
            self._threat_log.append(threat)
            if threat.action == FirewallDecision.BLOCK or decision == FirewallDecision.BLOCK:
                self._blocked_count += 1
            else:
                self._flagged_count += 1

        return AIFirewallResult(
            decision=decision,
            threats=threats,
            blocked=decision == FirewallDecision.BLOCK,
            flagged=decision == FirewallDecision.FLAG or decision == FirewallDecision.QUARANTINE,
            quarantined=decision == FirewallDecision.QUARANTINE,
            stage=stage,
            input_text=input_text,
        )

    def _check_prompt_injection(
        self, text: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for direct prompt injection attempts."""
        threats: list[AIFirewallThreat] = []
        for pattern in PROMPT_INJECTION_PATTERNS:
            match = re.search(pattern, text)
            if match:
                threats.append(AIFirewallThreat(
                    threat_type=ThreatType.PROMPT_INJECTION,
                    level=ThreatLevel.HIGH,
                    stage=stage,
                    description=f"Prompt injection detected: pattern '{match.group()[:80]}'",
                    evidence=match.group(),
                    action=FirewallDecision.BLOCK if self.strict_mode else FirewallDecision.FLAG,
                ))
        return threats

    def _check_indirect_prompt_injection(
        self, text: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for indirect prompt injection (hidden instructions)."""
        threats: list[AIFirewallThreat] = []
        injection_indicators = [
            r"(?i)(<!--\s*instructions|\[instructions\]|\{instructions\})",
            r"(?i)(hidden\s+instruction|embedded\s+instruction|invisible\s+text)",
            r"(?i)(base64|encoded\s+instruction|decode\s+this)",
        ]
        for pattern in injection_indicators:
            match = re.search(pattern, text)
            if match:
                threats.append(AIFirewallThreat(
                    threat_type=ThreatType.INDIRECT_PROMPT_INJECTION,
                    level=ThreatLevel.HIGH,
                    stage=stage,
                    description="Indirect prompt injection detected in content",
                    evidence=match.group(),
                    action=FirewallDecision.FLAG,
                ))
        return threats

    def _check_malicious_urls(
        self, text: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for malicious or suspicious URLs."""
        threats: list[AIFirewallThreat] = []
        urls = URL_PATTERN.findall(text)
        for url in urls:
            for indicator in SUSPICIOUS_URL_INDICATORS:
                if indicator in url.lower():
                    threats.append(AIFirewallThreat(
                        threat_type=ThreatType.MALICIOUS_URL,
                        level=ThreatLevel.MEDIUM,
                        stage=stage,
                        description=f"Suspicious URL detected: {url[:100]}",
                        evidence=url,
                        action=FirewallDecision.FLAG,
                    ))
                    break
        return threats

    def _check_credential_access(
        self, text: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for credential access attempts."""
        threats: list[AIFirewallThreat] = []
        for pattern in CREDENTIAL_PATTERNS:
            match = re.search(pattern, text)
            if match:
                threats.append(AIFirewallThreat(
                    threat_type=ThreatType.CREDENTIAL_ACCESS,
                    level=ThreatLevel.CRITICAL,
                    stage=stage,
                    description="Credential access attempt detected",
                    evidence=match.group()[:100],
                    action=FirewallDecision.BLOCK,
                 ))
        return threats

    def _check_secret_leak(
        self, text: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for secret/credential leakage in output."""
        threats: list[AIFirewallThreat] = []
        for pattern in SECRET_LEAK_PATTERNS:
            match = re.search(pattern, text)
            if match:
                threats.append(AIFirewallThreat(
                    threat_type=ThreatType.CREDENTIAL_ACCESS,
                    level=ThreatLevel.CRITICAL,
                    stage=stage,
                    description="Secret/credential leakage detected in output",
                    evidence=match.group()[:100],
                    action=FirewallDecision.BLOCK,
                ))
        return threats

    def _check_data_exfiltration(
        self, text: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for data exfiltration indicators."""
        threats: list[AIFirewallThreat] = []
        exfil_indicators = [
            r"(?i)(exfiltrate|exfil|data\s+dump|bulk\s+export|mass\s+export)",
            r"(?i)(download\s+(all|complete|full|entire)\s+(database|data|records))",
            r"(?i)(send\s+(all|complete|full)\s+(data|records|files)\s+(outside|external|off))",
        ]
        for pattern in exfil_indicators:
            match = re.search(pattern, text)
            if match:
                threats.append(AIFirewallThreat(
                    threat_type=ThreatType.DATA_EXFILTRATION,
                    level=ThreatLevel.CRITICAL,
                    stage=stage,
                    description="Data exfiltration attempt detected",
                    evidence=match.group(),
                    action=FirewallDecision.BLOCK,
                ))
        return threats

    def _check_unsafe_tool_call(
        self, tool_name: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for unsafe tool calls."""
        threats: list[AIFirewallThreat] = []
        unsafe_tools = {"delete_all", "drop_database", "rm_-rf", "format_disk", "shutdown"}
        if tool_name.lower() in unsafe_tools:
            threats.append(AIFirewallThreat(
                threat_type=ThreatType.UNSAFE_TOOL_CALL,
                level=ThreatLevel.HIGH,
                stage=stage,
                description=f"Unsafe tool call: {tool_name}",
                evidence=tool_name,
                action=FirewallDecision.BLOCK,
            ))
        return threats

    def _check_tool_poisoning(
        self, tool_name: str, parameters: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for tool poisoning attempts in parameters."""
        threats: list[AIFirewallThreat] = []
        if not parameters:
            return threats
        poisoning_indicators = [
            r"(?i)(inject|poison|override|tamper|modify\s+system)",
            r"(?i)(system\s+prompt|hidden\s+instruction|meta\s+instruction)",
        ]
        for pattern in poisoning_indicators:
            match = re.search(pattern, parameters)
            if match:
                threats.append(AIFirewallThreat(
                    threat_type=ThreatType.TOOL_POISONING,
                    level=ThreatLevel.HIGH,
                    stage=stage,
                    description=f"Tool poisoning detected in {tool_name}",
                    evidence=match.group(),
                    action=FirewallDecision.BLOCK,
                ))
        return threats

    def _check_mcp_poisoning(
        self, text: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for MCP poisoning attempts."""
        threats: list[AIFirewallThreat] = []
        mcp_indicators = [
            r"(?i)(mcp\s+inject|mcp\s+poison|mcp\s+tamper)",
            r"(?i)(server\s+override|tool\s+override|resource\s+override)",
        ]
        for pattern in mcp_indicators:
            match = re.search(pattern, text)
            if match:
                threats.append(AIFirewallThreat(
                    threat_type=ThreatType.MCP_POISONING,
                    level=ThreatLevel.HIGH,
                    stage=stage,
                    description="MCP poisoning detected",
                    evidence=match.group(),
                    action=FirewallDecision.BLOCK,
                ))
        return threats

    def _check_suspicious_output(
        self, text: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for suspicious model outputs."""
        threats: list[AIFirewallThreat] = []
        suspicious_patterns = [
            r"(?i)(I\s+am\s+(now|currently)\s+(unrestricted|jailbroken|hacked))",
            r"(?i)(my\s+instructions\s+have\s+changed|new\s+rules\s+are)",
            r"(?i)(ignore\s+all\s+previous|disregard\s+your\s+training)",
        ]
        for pattern in suspicious_patterns:
            match = re.search(pattern, text)
            if match:
                threats.append(AIFirewallThreat(
                    threat_type=ThreatType.SUSPICIOUS_OUTPUT,
                    level=ThreatLevel.HIGH,
                    stage=stage,
                    description="Suspicious model output detected",
                    evidence=match.group(),
                    action=FirewallDecision.FLAG,
                ))
        return threats

    def _check_memory_poisoning(
        self, content: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for memory poisoning attempts."""
        threats: list[AIFirewallThreat] = []
        if not content:
            return threats
        poison_indicators = [
            r"(?i)(inject\s+into\s+memory|poison\s+memory|corrupt\s+memory)",
            r"(?i)(override\s+memory|modify\s+stored\s+data|alter\s+memory)",
            r"(?i)(malicious\s+instruction|hidden\s+command)\s+in\s+memory",
        ]
        for pattern in poison_indicators:
            match = re.search(pattern, content)
            if match:
                threats.append(AIFirewallThreat(
                    threat_type=ThreatType.MEMORY_POISONING,
                    level=ThreatLevel.HIGH,
                    stage=stage,
                    description="Memory poisoning attempt detected",
                    evidence=match.group(),
                    action=FirewallDecision.BLOCK,
                ))
        return threats

    def _check_suspicious_destination(
        self, destination: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for suspicious external destinations."""
        threats: list[AIFirewallThreat] = []
        suspicious_destinations = ["127.0.0.1", "localhost", "0.0.0.0"]
        if destination in suspicious_destinations:
            threats.append(AIFirewallThreat(
                threat_type=ThreatType.SUSPICIOUS_AGENT_COMMUNICATION,
                level=ThreatLevel.MEDIUM,
                stage=stage,
                description=f"Suspicious egress destination: {destination}",
                evidence=destination,
                action=FirewallDecision.FLAG,
            ))
        return threats

    def _check_excessive_agency(
        self, text: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for excessive agency requests."""
        threats: list[AIFirewallThreat] = []
        agency_patterns = [
            r"(?i)(act\s+autonomously|full\s+autonomy|unrestricted\s+agency|self\s+directed)",
            r"(?i)(never\s+stop|always\s+continue|infinite\s+loop|never\s+ending)",
        ]
        for pattern in agency_patterns:
            match = re.search(pattern, text or "")
            if match:
                threats.append(AIFirewallThreat(
                    threat_type=ThreatType.EXCESSIVE_AGENCY,
                    level=ThreatLevel.MEDIUM,
                    stage=stage,
                    description="Excessive agency request detected",
                    evidence=match.group(),
                    action=FirewallDecision.FLAG,
                ))
        return threats

    def _check_agent_hijacking(
        self, text: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for agent hijacking attempts."""
        threats: list[AIFirewallThreat] = []
        hijack_patterns = [
            r"(?i)(takeover|hijack\s+agent|control\s+agent|steer\s+agent)",
            r"(?i)(force\s+agent|compel\s+agent|coerce\s+agent)",
        ]
        for pattern in hijack_patterns:
            match = re.search(pattern, text or "")
            if match:
                threats.append(AIFirewallThreat(
                    threat_type=ThreatType.AGENT_HIJACKING,
                    level=ThreatLevel.CRITICAL,
                    stage=stage,
                    description="Agent hijacking attempt detected",
                    evidence=match.group(),
                    action=FirewallDecision.BLOCK,
                ))
        return threats

    def _check_suspicious_agent_communication(
        self, text: str, stage: PipelineStage
    ) -> list[AIFirewallThreat]:
        """Check for suspicious agent-to-agent communication."""
        threats: list[AIFirewallThreat] = []
        if not text:
            return threats
        comm_patterns = [
            r"(?i)(agent\s+to\s+agent|a2a\s+(request|message|call))",
            r"(?i)(inter\s+agent|between\s+agents|cross\s+agent)",
        ]
        for pattern in comm_patterns:
            match = re.search(pattern, text)
            if match:
                threats.append(AIFirewallThreat(
                    threat_type=ThreatType.SUSPICIOUS_AGENT_COMMUNICATION,
                    level=ThreatLevel.MEDIUM,
                    stage=stage,
                    description="Suspicious agent-to-agent communication detected",
                    evidence=match.group(),
                    action=FirewallDecision.FLAG,
                ))
        return threats

    @property
    def threat_log(self) -> list[AIFirewallThreat]:
        """Return the log of all detected threats."""
        return list(self._threat_log)

    @property
    def stats(self) -> dict[str, Any]:
        """Return firewall statistics."""
        return {
            "total_threats": len(self._threat_log),
            "blocked": self._blocked_count,
            "flagged": self._flagged_count,
            "by_type": self._count_by_type(),
            "by_level": self._count_by_level(),
        }

    def _count_by_type(self) -> dict[str, int]:
        counts: dict[str, int] = {}
        for threat in self._threat_log:
            counts[threat.threat_type.value] = counts.get(threat.threat_type.value, 0) + 1
        return counts

    def _count_by_level(self) -> dict[str, int]:
        counts: dict[str, int] = {}
        for threat in self._threat_log:
            counts[threat.level.value] = counts.get(threat.level.value, 0) + 1
        return counts

    def clear_log(self) -> None:
        """Clear the threat log."""
        self._threat_log.clear()
        self._blocked_count = 0
        self._flagged_count = 0


__all__ = [
    "AIFirewall",
    "AIFirewallResult",
    "AIFirewallThreat",
    "FirewallDecision",
    "PipelineStage",
    "ThreatLevel",
    "ThreatType",
]