"""
PII Redaction Module

Automatically detects and redacts personally identifiable information (PII)
from SDK payloads before they are sent to the Sentrial API.
"""

import re
import hashlib
from dataclasses import dataclass
from typing import Any, List, Optional


# ============================================================================
# Types
# ============================================================================


@dataclass
class PiiCustomPattern:
    """A custom regex pattern with a human-readable label."""
    pattern: str
    label: str


@dataclass
class PiiBuiltinPatterns:
    """Toggle individual builtin pattern categories."""
    emails: bool = True
    phones: bool = True
    ssns: bool = True
    credit_cards: bool = True
    ip_addresses: bool = True


@dataclass
class PiiConfig:
    """Full PII redaction configuration."""
    enabled: bool = False
    fields: Optional[List[str]] = None
    mode: str = 'label'
    enhanced_detection: bool = False
    builtin_patterns: Optional[PiiBuiltinPatterns] = None
    custom_patterns: Optional[List[PiiCustomPattern]] = None

    def __post_init__(self) -> None:
        if self.fields is None:
            self.fields = [
                'user_input',
                'assistant_output',
                'userInput',
                'assistantOutput',
                'tool_input',
                'tool_output',
                'toolInput',
                'toolOutput',
                'metadata',
                'reasoning',
            ]
        if self.builtin_patterns is None:
            self.builtin_patterns = PiiBuiltinPatterns()
        if self.custom_patterns is None:
            self.custom_patterns = []


# ============================================================================
# Builtin Regex Patterns (compiled)
# ============================================================================

# Email addresses (RFC 5322 simplified)
_EMAIL_PATTERN = re.compile(r'[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}')

# US phone numbers: (###) ###-####, ###-###-####, ### ### ####, +1##########
_PHONE_PATTERN = re.compile(r'(?:\+?1[-.\s]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b')

# US Social Security Numbers: ###-##-####
_SSN_PATTERN = re.compile(r'\b\d{3}-\d{2}-\d{4}\b')

# Credit card numbers: 13-19 digits, optionally separated by spaces or dashes
_CREDIT_CARD_PATTERN = re.compile(r'\b(?:\d[-\s]?){13,19}\b')

# IPv4 addresses
_IP_ADDRESS_PATTERN = re.compile(
    r'\b(?:(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\b'
)

# Map of builtin pattern keys to their regex and label
_BUILTIN_PATTERNS = {
    'emails': {'pattern': _EMAIL_PATTERN, 'label': 'EMAIL'},
    'phones': {'pattern': _PHONE_PATTERN, 'label': 'PHONE'},
    'ssns': {'pattern': _SSN_PATTERN, 'label': 'SSN'},
    'credit_cards': {'pattern': _CREDIT_CARD_PATTERN, 'label': 'CREDIT_CARD'},
    'ip_addresses': {'pattern': _IP_ADDRESS_PATTERN, 'label': 'IP_ADDRESS'},
}


# ============================================================================
# Core Functions
# ============================================================================


def _hash_value(value: str) -> str:
    """
    Hash a value using SHA-256 and return the first 6 hex characters.
    Used in 'hash' mode to produce a short, deterministic pseudonym.
    """
    return hashlib.sha256(value.encode('utf-8')).hexdigest()[:6]


def _replace_match(match_str: str, label: str, mode: str) -> str:
    """
    Produce the replacement string for a matched PII value.

    - 'label' mode:  [LABEL]
    - 'hash' mode:   [PII:abcdef]
    - 'remove' mode: '' (empty string)
    """
    if mode == 'label':
        return f'[{label}]'
    elif mode == 'hash':
        return f'[PII:{_hash_value(match_str)}]'
    elif mode == 'remove':
        return ''
    return f'[{label}]'


def redact_string(
    value: str,
    mode: str,
    builtin: PiiBuiltinPatterns,
    custom: List[PiiCustomPattern],
) -> str:
    """
    Run all enabled builtin and custom patterns against a single string,
    replacing every match according to the specified mode.
    """
    result = value

    # Apply enabled builtin patterns
    for key, config in _BUILTIN_PATTERNS.items():
        if getattr(builtin, key, False):
            pattern = config['pattern']
            label = config['label']
            result = pattern.sub(
                lambda m, lbl=label: _replace_match(m.group(), lbl, mode),
                result,
            )

    # Apply custom patterns
    for cp in custom:
        compiled = re.compile(cp.pattern)
        result = compiled.sub(
            lambda m, lbl=cp.label: _replace_match(m.group(), lbl, mode),
            result,
        )

    return result


def redact_value(
    value: Any,
    mode: str,
    builtin: PiiBuiltinPatterns,
    custom: List[PiiCustomPattern],
) -> Any:
    """
    Recursively redact PII from a value. Handles strings, lists, dicts,
    and passes through primitives (numbers, booleans, None) unchanged.
    """
    if isinstance(value, str):
        return redact_string(value, mode, builtin, custom)

    if isinstance(value, list):
        return [redact_value(item, mode, builtin, custom) for item in value]

    if isinstance(value, dict):
        return {
            k: redact_value(v, mode, builtin, custom)
            for k, v in value.items()
        }

    # Primitives (int, float, bool, None) pass through unchanged
    return value


def redact_payload(
    payload: dict,
    config: PiiConfig,
) -> dict:
    """
    Main entry point: redact configured fields in a payload object.

    Only fields listed in the PiiConfig (or the defaults) are redacted.
    Other fields are returned unchanged.

    Args:
        payload: The data payload (e.g., session params, tool call params)
        config:  PII configuration from the client

    Returns:
        A new payload dict with PII redacted from the configured fields
    """
    if not config.enabled:
        return payload

    mode = config.mode or 'label'
    fields = config.fields or []
    builtin = config.builtin_patterns or PiiBuiltinPatterns()
    custom = config.custom_patterns or []

    result = dict(payload)

    for f in fields:
        if f in result and result[f] is not None:
            result[f] = redact_value(result[f], mode, builtin, custom)

    return result
