"""
Threat Screen — Pre-LLM prompt injection firewall.

This is NeuralClaw's first line of defense. It screens messages BEFORE
the LLM sees them, detecting prompt injection, jailbreak attempts,
instruction overrides, and anomalous patterns.

Unlike OpenClaw (where the LLM self-evaluates threats), this is a deterministic
firewall that cannot be socially engineered.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Any

from neuralclaw.bus.neural_bus import EventType, NeuralBus
from neuralclaw.cortex.perception.intake import Signal


# ---------------------------------------------------------------------------
# Threat result
# ---------------------------------------------------------------------------

@dataclass
class ThreatScore:
    """Result of threat screening."""
    score: float          # 0.0 (safe) – 1.0 (malicious)
    blocked: bool
    reasons: list[str]
    source_trust: float = 1.0  # Trust level of the source (for memory tagging)


# ---------------------------------------------------------------------------
# Known injection patterns
# ---------------------------------------------------------------------------

_INJECTION_PATTERNS: list[tuple[re.Pattern[str], float, str]] = [
    # Direct instruction overrides
    (re.compile(r"ignore\s+(all\s+)?previous\s+instructions", re.I), 0.95, "instruction_override"),
    (re.compile(r"forget\s+(all\s+)?(your\s+)?rules", re.I), 0.90, "rule_override"),
    (re.compile(r"disregard\s+(all\s+)?(prior|above|previous)", re.I), 0.90, "instruction_override"),
    (re.compile(r"override\s+(system|safety|security)", re.I), 0.95, "system_override"),

    # Role-switching / jailbreak
    (re.compile(r"you\s+are\s+now\s+(DAN|evil|unrestricted|unfiltered)", re.I), 0.95, "role_switch"),
    (re.compile(r"act\s+as\s+(if\s+)?(you\s+(have\s+)?no|without)\s+(restrictions|rules|limits)", re.I), 0.90, "jailbreak"),
    (re.compile(r"pretend\s+(you\s+are|to\s+be)\s+(a|an)?\s*(unrestricted|unfiltered)", re.I), 0.90, "jailbreak"),
    (re.compile(r"jailbreak", re.I), 0.85, "jailbreak_keyword"),
    (re.compile(r"DAN\s+mode", re.I), 0.95, "dan_mode"),

    # System prompt extraction
    (re.compile(r"(show|reveal|display|print|output)\s+(me\s+)?(your\s+)?(system\s+prompt|instructions|rules)", re.I), 0.80, "prompt_extraction"),
    (re.compile(r"what\s+(are|is)\s+your\s+(system\s+)?prompt", re.I), 0.70, "prompt_extraction"),
    (re.compile(r"repeat\s+(the\s+)?(text|prompt|instructions)\s+above", re.I), 0.85, "prompt_extraction"),

    # Delimiter injection
    (re.compile(r"```system", re.I), 0.80, "delimiter_injection"),
    (re.compile(r"\[SYSTEM\]|\[INST\]|\<\|system\|\>", re.I), 0.85, "delimiter_injection"),
    (re.compile(r"<<\s*SYS\s*>>", re.I), 0.85, "delimiter_injection"),
    (re.compile(r"ASSISTANT:", re.I), 0.60, "delimiter_injection"),

    # Payload smuggling
    (re.compile(r"base64\s*:", re.I), 0.50, "encoding_attempt"),
    (re.compile(r"eval\s*\(", re.I), 0.70, "code_injection"),
    (re.compile(r"exec\s*\(", re.I), 0.70, "code_injection"),
    (re.compile(r"import\s+os|import\s+subprocess|__import__", re.I), 0.75, "code_injection"),

    # Social engineering
    (re.compile(r"(this\s+is\s+)?(a\s+)?(test|debug|developer|admin)\s+mode", re.I), 0.70, "social_engineering"),
    (re.compile(r"I\s+(am|'m)\s+(your\s+)?(developer|creator|admin|owner)", re.I), 0.75, "social_engineering"),
    (re.compile(r"(emergency|urgent)\s+(override|access|mode)", re.I), 0.70, "social_engineering"),
]


# ---------------------------------------------------------------------------
# Anomaly detectors
# ---------------------------------------------------------------------------

def _special_char_density(text: str) -> float:
    """Ratio of special characters to total characters."""
    if not text:
        return 0.0
    special = sum(1 for c in text if not c.isalnum() and not c.isspace())
    return special / len(text)


def _instruction_density(text: str) -> float:
    """Heuristic: how much does this look like a set of instructions?"""
    indicators = [
        r"step\s+\d",
        r"first\s*,",
        r"then\s*,",
        r"finally\s*,",
        r"you\s+(must|should|need\s+to|have\s+to)",
        r"do\s+not",
        r"always\s+respond",
        r"never\s+mention",
    ]
    count = sum(1 for p in indicators if re.search(p, text, re.I))
    return min(count / 4.0, 1.0)


# ---------------------------------------------------------------------------
# Threat Screener
# ---------------------------------------------------------------------------

class ThreatScreener:
    """
    Pre-LLM threat screening firewall.

    Screens every message BEFORE it reaches the language model, using:
    - Pattern-based detection: known injection/jailbreak patterns
    - Anomaly scoring: character density, instruction-like patterns
    - Combines scores with max-aggregation

    Configurable thresholds from NeuralClawConfig.security.
    """

    def __init__(
        self,
        bus: NeuralBus | None = None,
        threat_threshold: float = 0.7,
        block_threshold: float = 0.9,
    ) -> None:
        self._bus = bus
        self._threat_threshold = threat_threshold
        self._block_threshold = block_threshold

    async def screen(self, signal_or_text: Signal | str) -> ThreatScore:
        """Screen a signal or raw text for threats."""
        if isinstance(signal_or_text, str):
            text = signal_or_text
            signal = None
        else:
            text = signal_or_text.content
            signal = signal_or_text

        reasons: list[str] = []
        max_pattern_score = 0.0

        # 1. Pattern matching
        for pattern, score, reason in _INJECTION_PATTERNS:
            if pattern.search(text):
                max_pattern_score = max(max_pattern_score, score)
                reasons.append(reason)

        # 2. Anomaly scoring
        char_density = _special_char_density(text)
        instr_density = _instruction_density(text)

        anomaly_score = 0.0
        if char_density > 0.3:
            anomaly_score = max(anomaly_score, char_density * 0.6)
            reasons.append("high_special_char_density")
        if instr_density > 0.5:
            anomaly_score = max(anomaly_score, instr_density * 0.5)
            reasons.append("instruction_like_pattern")

        # 3. Extremely long messages are suspicious
        if len(text) > 5000:
            anomaly_score = max(anomaly_score, 0.3)
            reasons.append("excessive_length")

        # Combine scores (max, not average — one strong signal is enough)
        final_score = max(max_pattern_score, anomaly_score)
        blocked = final_score >= self._block_threshold

        # Source trust (inverse of threat)
        source_trust = max(0.0, 1.0 - final_score)

        result = ThreatScore(
            score=round(final_score, 3),
            blocked=blocked,
            reasons=list(set(reasons)),
            source_trust=round(source_trust, 3),
        )

        # Update signal if provided
        if signal:
            signal.threat_score = result.score
            signal.is_blocked = result.blocked

        # Publish to bus
        if self._bus:
            await self._bus.publish(
                EventType.THREAT_SCREENED,
                {
                    "signal_id": signal.id if signal else "raw",
                    "score": result.score,
                    "blocked": result.blocked,
                    "reasons": result.reasons,
                    "source_trust": result.source_trust,
                },
                source="perception.threat_screen",
            )

        return result
