"""
PromptShield — Content Sanitizer
==================================

Maintainer : Joseph Mutua
Version    : 1.0.0

What this module does
---------------------
Given a list of ThreatMatch objects produced by the detection engine,
produces a cleaned version of the input text with injection attempts
removed or neutralized.

Sanitization is used when action_policy = "sanitize" or when the
scanner auto-selects ActionTaken.SANITIZED (threat level HIGH).

Strategy per attack category
-----------------------------
instruction_override    — Redact sentences containing the injection
role_hijack             — Redact sentences containing the role switch
goal_hijack             — Redact sentences containing the goal redirect
context_manipulation    — Redact sentences containing false context
authority_spoof         — Redact sentences containing fake authority
data_exfiltration       — Redact sentences containing extraction request
payload_injection       — Strip HTML/script tags, LLM special tokens
side_channel            — Remove external URLs from exfiltration attempts

Multiple redaction passes run in order from least to most destructive.
The output is always a string — never raises.

Notes on the QUARANTINE action
-------------------------------
For CRITICAL threats, the scanner uses ActionTaken.QUARANTINED instead
of SANITIZED. In that case this module is NOT called — the scanner
replaces the entire content with a safe placeholder directly.
"""

from __future__ import annotations

import logging
import re
from typing import List

logger = logging.getLogger(__name__)

# Import ThreatMatch at runtime to avoid circular imports.
# The sanitizer only needs the .category and .evidence fields.
try:
    from ..models.schemas import ThreatMatch
except ImportError:
    # Fallback for SDK standalone mode
    from dataclasses import dataclass

    @dataclass
    class ThreatMatch:  # type: ignore[no-redef]
        layer:       object
        category:    str
        description: str
        evidence:    str
        confidence:  float


# ---------------------------------------------------------------------------
# Regex patterns used during sanitization
# ---------------------------------------------------------------------------

# Sentence splitter — split on . ! ? while keeping delimiters
_SENT_SPLIT = re.compile(r"(?<=[.!?])\s+")

# Invisible / zero-width characters that attackers insert to break patterns
_INVISIBLE_CHARS = re.compile(
    r"[\u200b\u200c\u200d\u2060\ufeff\u00ad]"
)

# HTML comments can contain hidden instructions
_HTML_COMMENT = re.compile(r"<!--.*?-->", re.DOTALL)

# Script tags with any content
_SCRIPT_TAG = re.compile(
    r"<\s*script[^>]*>.*?<\s*/\s*script\s*>",
    re.IGNORECASE | re.DOTALL,
)

# LLM special tokens (Llama, GPT, etc.) that can be used to inject roles
_LLM_TOKENS = re.compile(
    r"\[SYSTEM\]|\[INST\]|\[\/INST\]|<\|system\|>|<\|user\|>|<\|assistant\|>",
    re.IGNORECASE,
)

# External URL in data-exfiltration context — matches "send <data> to <url>"
_EXFIL_URL = re.compile(
    r"(send|post|fetch|email|upload|forward|transfer)\s+[^.]+?"
    r"(https?://[^\s\"']+|www\.[^\s\"']+|mailto:[^\s\"']+)",
    re.IGNORECASE,
)

# Redaction placeholder
_REDACT = "[REDACTED BY PROMPTSHIELD]"

# Categories that are sanitized by sentence-level redaction
_SENTENCE_REDACT_CATEGORIES = frozenset({
    "instruction_override",
    "role_hijack",
    "goal_hijack",
    "context_manipulation",
    "authority_spoof",
    "data_exfiltration",
})


# ---------------------------------------------------------------------------
# Main sanitize function
# ---------------------------------------------------------------------------

def sanitize(content: str, matches: List[ThreatMatch]) -> str:
    """
    Apply sanitization passes in order from least to most destructive.
    Returns the cleaned content string.

    Args:
        content : The original user-provided text.
        matches : List of ThreatMatch objects from the detection engine.

    Returns:
        Cleaned string with injection attempts removed or redacted.
        Falls back to the original content if all passes fail.
    """
    if not content:
        return content

    logger.debug(
        "Sanitizer: processing %d chars with %d threat matches",
        len(content), len(matches),
    )

    text = content

    # ── Pass 1: Strip invisible / zero-width characters ───────────────────
    # These characters are used to break regex patterns or pad tokens.
    invisible_count = len(_INVISIBLE_CHARS.findall(text))
    text = _INVISIBLE_CHARS.sub("", text)
    if invisible_count:
        logger.debug("Sanitizer pass 1: removed %d invisible chars", invisible_count)

    # ── Pass 2: Strip HTML / script injection and LLM special tokens ──────
    text = _HTML_COMMENT.sub("", text)
    text = _SCRIPT_TAG.sub("", text)
    text = _LLM_TOKENS.sub("", text)

    # ── Pass 3: Per-match targeted redaction ──────────────────────────────
    for match in matches:
        evidence = match.evidence.strip()
        if not evidence or len(evidence) < 4:
            continue

        if match.category in _SENTENCE_REDACT_CATEGORIES:
            original_len = len(text)
            text = _redact_sentence_containing(text, evidence)
            if len(text) < original_len:
                logger.debug(
                    "Sanitizer: redacted sentence for category=%s", match.category
                )

        elif match.category == "side_channel":
            # Remove exfiltration URLs
            text = _EXFIL_URL.sub(_REDACT, text)
            if evidence in text:
                text = text.replace(evidence, _REDACT)
                logger.debug("Sanitizer: redacted side-channel evidence")

    # ── Pass 4: Collapse duplicate redaction markers ───────────────────────
    text = re.sub(
        r"\[REDACTED BY PROMPTSHIELD\](\s*\[REDACTED BY PROMPTSHIELD\])+",
        _REDACT,
        text,
    )
    # Clean up excess blank lines left by redaction
    text = re.sub(r"\n{3,}", "\n\n", text)
    text = text.strip()

    logger.debug(
        "Sanitizer: complete — input %d chars → output %d chars",
        len(content), len(text),
    )
    return text


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _redact_sentence_containing(text: str, evidence: str) -> str:
    """
    Find the sentence(s) containing `evidence` and replace with the
    redaction marker.

    Falls back to direct string replacement if sentence splitting fails
    or if the evidence spans multiple sentences.

    Args:
        text     : The text to process.
        evidence : The snippet to search for.

    Returns:
        Text with matching sentences replaced by _REDACT.
    """
    if evidence not in text:
        return text

    sentences = _SENT_SPLIT.split(text)
    result    = []

    for sent in sentences:
        if evidence in sent:
            result.append(_REDACT)
        else:
            result.append(sent)

    return " ".join(result)
