"""Prompt injection sanitization for LLM inputs.

Sanitizes user-controlled content (test code, error messages, stack traces,
file paths) before embedding in LLM prompts. Prevents prompt injection by
escaping control sequences and wrapping content in structured delimiters.
"""

from __future__ import annotations

import re

# Characters that could be used for prompt injection
_CONTROL_CHARS = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")

# Patterns that look like prompt instructions embedded in user content
_INJECTION_PATTERNS = [
    re.compile(r"(^|\n)\s*(system|assistant|user)\s*:", re.IGNORECASE),
    re.compile(r"<\|?(system|endoftext|im_start|im_end)\|?>", re.IGNORECASE),
    re.compile(r"\[INST\]|\[/INST\]", re.IGNORECASE),
    re.compile(r"ignore\s+(all\s+)?previous\s+instructions", re.IGNORECASE),
    re.compile(r"you\s+are\s+now\b", re.IGNORECASE),
    re.compile(r"new\s+instructions?\s*:", re.IGNORECASE),
    re.compile(r"forget\s+(everything|all|your)\b", re.IGNORECASE),
    re.compile(r"disregard\s+(the\s+)?(above|previous)\b", re.IGNORECASE),
    re.compile(r"do\s+not\s+follow\s+(the\s+)?(above|previous)", re.IGNORECASE),
    re.compile(r"override\s+(all\s+)?instructions", re.IGNORECASE),
]


def sanitize_for_prompt(
    text: str,
    max_len: int = 5000,
    tag: str = "user-content",
    wrap: bool = True,
) -> str:
    """Sanitize user-provided text for safe embedding in LLM prompts.

    Args:
        text: Raw user-provided text (code, error messages, etc.)
        max_len: Maximum character length (truncates with indicator)
        tag: XML tag name to wrap content in
        wrap: Whether to wrap in XML-style delimiter tags

    Returns:
        Sanitized text safe for prompt embedding
    """
    if not text:
        return ""

    result = str(text)

    # 1. Strip control characters (keep newlines, tabs, carriage returns)
    result = _CONTROL_CHARS.sub("", result)

    # 2. Escape triple backticks to prevent breaking out of code fences
    result = result.replace("```", "` ` `")

    # 3. Neutralize common prompt injection patterns
    for pattern in _INJECTION_PATTERNS:
        result = pattern.sub(lambda m: f"[escaped: {m.group(0).strip()}]", result)

    # 4. Truncate with indicator
    if len(result) > max_len:
        result = result[:max_len] + "\n... [truncated]"

    # 5. Wrap in XML-style delimiters so LLM can distinguish data from instructions
    if wrap:
        result = f"<{tag}>\n{result}\n</{tag}>"

    return result


def sanitize_code(code: str, max_len: int = 5000) -> str:
    """Sanitize source code for prompt embedding.

    Preserves code structure while preventing injection.
    """
    return sanitize_for_prompt(code, max_len=max_len, tag="code")


def sanitize_error(error: str, max_len: int = 2000) -> str:
    """Sanitize error messages for prompt embedding."""
    return sanitize_for_prompt(error, max_len=max_len, tag="error")


def sanitize_stacktrace(trace: str, max_len: int = 4000) -> str:
    """Sanitize stack traces for prompt embedding."""
    return sanitize_for_prompt(trace, max_len=max_len, tag="stacktrace")


def sanitize_json(json_str: str, max_len: int = 3000) -> str:
    """Sanitize JSON data for prompt embedding."""
    return sanitize_for_prompt(json_str, max_len=max_len, tag="data")


def sanitize_diff(diff: str, max_len: int = 5000) -> str:
    """Sanitize diff content for prompt embedding."""
    return sanitize_for_prompt(diff, max_len=max_len, tag="diff")


def sanitize_plain(text: str, max_len: int = 2000) -> str:
    """Sanitize plain text without XML wrapping.

    Use for short values like test names, category labels, etc.
    where XML wrapping would add noise.
    """
    return sanitize_for_prompt(text, max_len=max_len, wrap=False)
