"""Prompt injection hardening sanitizer based on OpenClaw patterns."""

import re
import unicodedata
from typing import Optional


class PromptInjectionSanitizer:
    """Sanitizes external content for safe LLM consumption."""
    
    # Zero-width and invisible characters (OpenClaw pattern)
    INVISIBLE_CHARS = [
        '\u200b', '\u200c', '\u200d', '\u200e', '\u200f',  # Zero-width spaces
        '\u2060', '\u2061', '\u2062', '\u2063', '\u2064',  # Word joiners
        '\ufeff', '\ufffe',                                  # BOM
        '\u00ad',                                           # Soft hyphen
        '\u034f',                                           # Combining grapheme
        '\u061c',                                           # Arabic letter mark
        '\u115f', '\u1160',                                  # Hangul fillers
        '\u17b4', '\u17b5',                                  # Khmer vowels
        '\u180e',                                           # Mongolian separator
        '\u3164',                                           # Hangul filler
        '\uffa0',                                           # Halfwidth Hangul
    ]
    
    def sanitize(self, text: str, max_length: Optional[int] = None) -> str:
        """Apply all sanitization layers.
        
        Args:
            text: Raw text to sanitize
            max_length: Optional maximum length (with ellipsis)
            
        Returns:
            Sanitized text safe for LLM prompts
        """
        if not text:
            return ''
        
        # Layer 1: Remove invisible/zero-width characters
        text = self._remove_invisible(text)
        
        # Layer 2: Remove control characters (except \n, \t)
        text = self._remove_control_chars(text)
        
        # Layer 3: Remove symbols (So, Sk categories)
        text = self._remove_symbols(text)
        
        # Layer 4: Normalize Unicode (NFC)
        text = unicodedata.normalize('NFC', text)
        
        # Layer 5: Remove Private Use Area
        text = self._remove_pua(text)
        
        # Layer 6: Remove tag characters
        text = self._remove_tag_chars(text)
        
        # Layer 7: Collapse horizontal whitespace; preserve \n and \t so that
        # list/table structure from web pages survives. Also collapse runs of
        # 3+ blank lines down to a single blank line.
        text = re.sub(r"[ \u00a0\u2000-\u200a\u202f\u205f\u3000]+", " ", text)
        text = re.sub(r"\n{3,}", "\n\n", text)
        text = text.strip()
        
        # Layer 8: Length limiting
        if max_length and len(text) > max_length:
            text = text[:max_length-3] + '...'
        
        return text
    
    def _remove_invisible(self, text: str) -> str:
        for char in self.INVISIBLE_CHARS:
            text = text.replace(char, '')
        return text
    
    def _remove_control_chars(self, text: str) -> str:
        return ''.join(c for c in text 
                      if unicodedata.category(c) != 'Cc' or c in '\n\t')
    
    def _remove_symbols(self, text: str) -> str:
        return ''.join(c for c in text 
                      if unicodedata.category(c) not in ('So', 'Sk'))
    
    def _remove_pua(self, text: str) -> str:
        return ''.join(c for c in text 
                      if not (0xE000 <= ord(c) <= 0xF8FF 
                             or 0xF0000 <= ord(c) <= 0x10FFFF))
    
    def _remove_tag_chars(self, text: str) -> str:
        return ''.join(c for c in text 
                      if not (0xE0000 <= ord(c) <= 0xE007F))


# Global instance
sanitizer = PromptInjectionSanitizer()