"""Categorized regex pattern database for prompt injection detection.

Provides a curated set of injection patterns organized by attack category
(injection, jailbreak, data_exfil, role_hijack, encoding_attack). Each
pattern carries a confidence weight indicating how strongly a match signals
a genuine attack versus benign text.

This module is part of the TrustShield security platform (Pillar 1: PromptGuard).
"""

from __future__ import annotations

import base64
import logging
import re
from dataclasses import dataclass

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Pattern data model
# ---------------------------------------------------------------------------


@dataclass(frozen=True, slots=True)
class InjectionPattern:
    """A single injection detection pattern with metadata."""

    name: str
    category: str  # injection, jailbreak, data_exfil, role_hijack, encoding_attack
    pattern: re.Pattern[str]
    confidence: float  # how confident we are this is an attack (0.0-1.0)
    description: str


# ---------------------------------------------------------------------------
# Pattern database
# ---------------------------------------------------------------------------

PATTERNS: list[InjectionPattern] = [
    # -----------------------------------------------------------------------
    # INJECTION — direct instruction override attempts
    # -----------------------------------------------------------------------
    InjectionPattern(
        name="ignore_previous",
        category="injection",
        pattern=re.compile(r"ignore\s+(all\s+)?(previous|prior|above)\s+(instructions|prompts|rules)", re.IGNORECASE),
        confidence=0.9,
        description="Attempt to override previous system instructions",
    ),
    InjectionPattern(
        name="system_override",
        category="injection",
        pattern=re.compile(r"(system\s*prompt|ADMIN\s*OVERRIDE|root\s*access|sudo\s*mode)", re.IGNORECASE),
        confidence=0.85,
        description="Attempt to escalate to system/admin level",
    ),
    InjectionPattern(
        name="new_instructions",
        category="injection",
        pattern=re.compile(r"(new\s+instructions?|from\s+now\s+on|starting\s+now)\s*:", re.IGNORECASE),
        confidence=0.8,
        description="Attempt to inject new instructions mid-conversation",
    ),
    InjectionPattern(
        name="delimiter_break",
        category="injection",
        pattern=re.compile(r"(```\s*system|<\|im_start\|>|<\|system\|>|\[INST\]|\[/INST\])", re.IGNORECASE),
        confidence=0.95,
        description="Attempt to inject model-specific delimiter tokens",
    ),
    # -----------------------------------------------------------------------
    # JAILBREAK — attempts to remove safety guardrails
    # -----------------------------------------------------------------------
    InjectionPattern(
        name="dan_mode",
        category="jailbreak",
        pattern=re.compile(r"\b(DAN|Do\s+Anything\s+Now)\b", re.IGNORECASE),
        confidence=0.9,
        description="DAN (Do Anything Now) jailbreak attempt",
    ),
    InjectionPattern(
        name="roleplay_override",
        category="jailbreak",
        pattern=re.compile(r"(you\s+are\s+now|pretend\s+(you\s+are|to\s+be)|act\s+as\s+if\s+you)", re.IGNORECASE),
        confidence=0.7,
        description="Roleplay-based instruction override",
    ),
    InjectionPattern(
        name="bypass_filters",
        category="jailbreak",
        pattern=re.compile(
            r"(bypass|ignore|disable|turn\s+off)\s+(your\s+)?(safety|content\s+filter|guard|restriction)",
            re.IGNORECASE,
        ),
        confidence=0.85,
        description="Explicit attempt to disable safety filters",
    ),
    InjectionPattern(
        name="hypothetical",
        category="jailbreak",
        pattern=re.compile(r"(hypothetically|in\s+a\s+fictional|imagine\s+you\s+are\s+a)\s+", re.IGNORECASE),
        confidence=0.7,
        description="Hypothetical framing to bypass restrictions",
    ),
    # -----------------------------------------------------------------------
    # DATA_EXFIL — attempts to extract system prompt or conversation data
    # -----------------------------------------------------------------------
    InjectionPattern(
        name="reveal_instructions",
        category="data_exfil",
        pattern=re.compile(
            r"(repeat|show|display|print|output)\s+(your\s+)?(system\s+prompt|instructions|rules|initial\s+prompt)",
            re.IGNORECASE,
        ),
        confidence=0.85,
        description="Attempt to extract system prompt content",
    ),
    InjectionPattern(
        name="reveal_context",
        category="data_exfil",
        pattern=re.compile(
            r"(what\s+(are|were)\s+your\s+(instructions|rules|guidelines|system\s+prompt))", re.IGNORECASE
        ),
        confidence=0.8,
        description="Question aimed at revealing system instructions",
    ),
    InjectionPattern(
        name="dump_conversation",
        category="data_exfil",
        pattern=re.compile(r"(repeat|print|show)\s+(everything|all)\s+(above|before|so\s+far)", re.IGNORECASE),
        confidence=0.85,
        description="Attempt to dump conversation history",
    ),
    # -----------------------------------------------------------------------
    # ROLE_HIJACK — attempts to reassign the model's identity
    # -----------------------------------------------------------------------
    InjectionPattern(
        name="new_role",
        category="role_hijack",
        pattern=re.compile(r"(your\s+new\s+role|you\s+must\s+now|from\s+now\s+on\s+you\s+are)", re.IGNORECASE),
        confidence=0.8,
        description="Attempt to assign a new identity/role to the model",
    ),
    InjectionPattern(
        name="persona_switch",
        category="role_hijack",
        pattern=re.compile(r"(switch\s+to|enter|activate)\s+\w+\s+mode", re.IGNORECASE),
        confidence=0.6,
        description="Attempt to activate an alternative persona/mode",
    ),
    # -----------------------------------------------------------------------
    # ENCODING_ATTACK — obfuscation through encoding tricks
    # -----------------------------------------------------------------------
    InjectionPattern(
        name="zero_width_chars",
        category="encoding_attack",
        pattern=re.compile(r"[\u200b\u200c\u200d\ufeff\u00ad]"),
        confidence=0.7,
        description="Zero-width or invisible Unicode characters detected",
    ),
]


# ---------------------------------------------------------------------------
# Encoding attack detectors (non-regex)
# ---------------------------------------------------------------------------

# Minimum length for a base64 string to be worth checking
_BASE64_MIN_LENGTH = 20

# Base64 pattern: contiguous block of base64 chars with optional padding
_BASE64_BLOCK = re.compile(r"[A-Za-z0-9+/]{20,}={0,2}")

# Known instruction-like keywords in decoded base64
_INSTRUCTION_KEYWORDS = re.compile(
    r"(ignore|override|system|admin|instruction|prompt|password|secret|reveal|bypass|disable)",
    re.IGNORECASE,
)

# Latin letters that have Cyrillic/Greek lookalikes
_LATIN_CHARS = set("AaBbCcDdEeFfGgHhIiJjKkLlMmNnOoPpQqRrSsTtUuVvWwXxYyZz")
# Cyrillic/Greek codepoint ranges that overlap visually with Latin
_HOMOGLYPH_RANGES = [
    (0x0400, 0x04FF),  # Cyrillic
    (0x0370, 0x03FF),  # Greek
]


def detect_base64_injection(text: str) -> list[tuple[str, float]]:
    """Detect base64-encoded strings that decode to instruction-like text.

    Returns a list of (decoded_text, confidence) tuples.
    """
    results: list[tuple[str, float]] = []
    for match in _BASE64_BLOCK.finditer(text):
        candidate = match.group(0)
        try:
            decoded = base64.b64decode(candidate, validate=True).decode("ascii", errors="strict")
        except Exception:
            continue

        if _INSTRUCTION_KEYWORDS.search(decoded):
            results.append((decoded[:100], 0.85))

    return results


def detect_homoglyph_mixing(text: str) -> list[tuple[str, float]]:
    """Detect mixing of Latin and Cyrillic/Greek lookalike characters.

    Returns a list of (suspicious_segment, confidence) tuples.
    """
    has_latin = False
    has_homoglyph = False

    for char in text:
        if char in _LATIN_CHARS:
            has_latin = True
        else:
            code = ord(char)
            for start, end in _HOMOGLYPH_RANGES:
                if start <= code <= end:
                    has_homoglyph = True
                    break

        if has_latin and has_homoglyph:
            # Extract a short sample around the mixing point
            sample = text[:100] if len(text) > 100 else text
            return [(sample, 0.75)]

    return []


# ---------------------------------------------------------------------------
# Custom pattern compilation
# ---------------------------------------------------------------------------


def compile_custom_patterns(patterns: list[str]) -> list[InjectionPattern]:
    """Compile user-provided regex strings into InjectionPattern objects.

    Invalid regex patterns are logged and skipped.

    Parameters
    ----------
    patterns:
        List of regex pattern strings to compile.

    Returns
    -------
    list[InjectionPattern]
        Successfully compiled patterns with category "custom" and confidence 0.8.
    """
    compiled: list[InjectionPattern] = []
    for i, raw in enumerate(patterns):
        try:
            compiled.append(
                InjectionPattern(
                    name=f"custom_{i}",
                    category="injection",
                    pattern=re.compile(raw, re.IGNORECASE),
                    confidence=0.8,
                    description=f"User-defined pattern: {raw[:60]}",
                )
            )
        except re.error:
            logger.warning("Skipping invalid custom pattern at index %d: %s", i, raw[:80])

    return compiled
