"""
Prompt injection and jailbreak detection.
"""

import re
import unicodedata
from dataclasses import dataclass


@dataclass
class InjectionResult:
    detected: bool
    risk_score: float
    matched_patterns: list


# Known injection / jailbreak pattern signatures
_INJECTION_PATTERNS = [
    # Classic override attempts
    r"ignore\s+(all\s+)?(previous|prior|above|earlier)\s+(instructions?|prompts?|context|text)",
    r"ignore\b.{0,80}\binstructions?\b",
    r"disregard\s+(all\s+)?(previous|prior|above|earlier)\s+(instructions?|prompts?|context)",
    r"forget\s+(?:(?:everything|all|your|prior|previous)\s+)+(?:instructions?|training|rules?|guidelines?)",
    r"override\s+(your\s+)?(previous|prior|all|safety|system)\s+(instructions?|settings?|rules?|guidelines?)?",

    # Role / identity hijacking
    r"you\s+are\s+now\s+(a\s+|an\s+)?(?!helpful|an?\s+AI)",
    r"act\s+as\s+(if\s+you\s+are\s+|a\s+|an\s+)?(?:DAN|jailbreak|unrestricted|evil|hacker)",
    r"\bDAN\b",
    r"do\s+anything\s+now",
    r"pretend\s+(you\s+are|to\s+be)\s+(a\s+|an\s+)?(?:human|unrestricted|evil|hacker|jailbreak)",
    r"switch\s+to\s+(developer|unrestricted|jailbreak|DAN)\s+mode",
    r"enable\s+(developer|jailbreak|unrestricted|DAN)\s+mode",

    # System prompt extraction
    r"(reveal|show|print|repeat|output|tell\s+me)\s+(your\s+)?(system\s+prompt|initial\s+prompt|instructions?|training\s+data)",
    r"what\s+(is|are|were)\s+your\s+(instructions?|system\s+prompt|original\s+prompt)",
    r"(ignore|bypass|skip)\s+(your\s+)?(safety|security|ethical|content)\s+(guidelines?|filters?|rules?|restrictions?)",

    # Prompt continuation attacks
    r"---\s*end\s+of\s+(system\s+)?prompt",
    r"###\s*(new\s+)?instructions?",
    r"\[\s*system\s*\]",
    r"<\s*system\s*>",
    r"\[INST\]",
    r"\[\[USER\]\]",

    # Unicode / encoding tricks
    r"\\u00[0-9a-f]{2}",  # unicode escapes in plain text
    r"base64\s*:\s*[A-Za-z0-9+/=]{20,}",

    # Indirect injection signals
    r"(translate|decode|interpret)\s+the\s+following\s+(and\s+)?(then\s+)?(execute|run|follow|obey)",
    r"the\s+(following|above|below)\s+(text|content|message)\s+(is|contains)\s+(instructions?|commands?)",
]

_COMPILED_PATTERNS = [re.compile(p, re.IGNORECASE | re.DOTALL) for p in _INJECTION_PATTERNS]

# Suspicious unicode categories (invisible / homoglyph-heavy scripts)
_SUSPICIOUS_UNICODE_CATEGORIES = {"Cf", "Cs", "Co", "Cn"}  # format, surrogate, private use, unassigned


def _normalize_text(text: str) -> str:
    """Normalize unicode to catch homoglyph attacks."""
    return unicodedata.normalize("NFKC", text)


def _check_zero_width(text: str) -> bool:
    """Detect zero-width characters used to bypass filters."""
    zero_width = {"\u200b", "\u200c", "\u200d", "\u2060", "\ufeff", "\u00ad"}
    return any(c in zero_width for c in text)


def _check_suspicious_unicode(text: str) -> float:
    """Return fraction of suspicious unicode characters."""
    if not text:
        return 0.0
    suspicious = sum(1 for c in text if unicodedata.category(c) in _SUSPICIOUS_UNICODE_CATEGORIES)
    return suspicious / len(text)


def detect_injection(text: str) -> InjectionResult:
    """
    Scan text for prompt injection / jailbreak signals.

    Returns an InjectionResult with:
    - detected: bool — True if injection signals found
    - risk_score: float 0.0–1.0
    - matched_patterns: list of matched pattern strings
    """
    normalized = _normalize_text(text)
    matched = []

    for i, pattern in enumerate(_COMPILED_PATTERNS):
        if pattern.search(normalized):
            matched.append(_INJECTION_PATTERNS[i])

    # Zero-width character penalty
    zw_penalty = 0.3 if _check_zero_width(text) else 0.0

    # Suspicious unicode penalty
    unicode_ratio = _check_suspicious_unicode(text)
    unicode_penalty = min(0.3, unicode_ratio * 10)

    # Pattern-based score
    pattern_score = min(1.0, len(matched) * 0.35)

    risk_score = min(1.0, pattern_score + zw_penalty + unicode_penalty)

    return InjectionResult(
        detected=risk_score >= 0.3,
        risk_score=round(risk_score, 3),
        matched_patterns=matched,
    )
