"""
PII detection, redaction, and rehydration.
"""

import re
from dataclasses import dataclass, field
from typing import Optional

# -------------------------------------------------------------------
# PII regex patterns  (ordered: most specific first)
# -------------------------------------------------------------------

_PII_PATTERNS = [
    # SSN  (US Social Security)
    ("SSN", re.compile(r"\b(?!000|666|9\d{2})\d{3}[- ](?!00)\d{2}[- ](?!0000)\d{4}\b")),
    # Credit card  (Visa, MC, Amex, Discover — with or without separators)
    ("CREDIT_CARD", re.compile(
        r"\b(?:4[0-9]{12}(?:[0-9]{3})?|5[1-5][0-9]{14}|3[47][0-9]{13}|"
        r"6(?:011|5[0-9]{2})[0-9]{12})[- ]?(?:[0-9]{4}[- ]?){0,3}\b"
    )),
    # Email
    ("EMAIL", re.compile(r"\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}\b")),
    # Phone (US-centric but catches most formats)
    ("PHONE", re.compile(
        r"\b(?:\+?1[-.\s]?)?(?:\(?\d{3}\)?[-.\s]?)?\d{3}[-.\s]?\d{4}\b"
    )),
    # IPv4
    ("IP_ADDRESS", 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"
    )),
    # Driver's license (basic US format — letter(s) + digits)
    ("DRIVERS_LICENSE", re.compile(r"\b[A-Z]{1,2}\d{5,8}\b")),
    # Date of birth patterns
    ("DATE_OF_BIRTH", re.compile(
        r"\b(?:DOB|date\s+of\s+birth|born\s+on)[:\s]+\d{1,2}[\/\-]\d{1,2}[\/\-]\d{2,4}\b",
        re.IGNORECASE
    )),
    # Passport (US format — letter + 8 digits)
    ("PASSPORT", re.compile(r"\b[A-Z][0-9]{8}\b")),
    # URL (catch potentially sensitive endpoints)
    ("URL", re.compile(r"https?://[^\s\"'<>\])]+")),
]


@dataclass
class RedactionResult:
    redacted_text: str
    mapping: dict = field(default_factory=dict)  # placeholder → original value
    entity_counts: dict = field(default_factory=dict)  # entity_type → count


def redact(text: str, entities: Optional[list] = None) -> RedactionResult:
    """
    Redact PII from text, replacing each match with a typed placeholder.

    Args:
        text: Input text to redact.
        entities: Optional list of entity type strings to target (e.g. ["EMAIL", "PHONE"]).
                  If None, all entity types are checked.

    Returns:
        RedactionResult with .redacted_text and .mapping dict.
        mapping maps placeholder → original value for rehydration.
    """
    if entities is not None:
        entities_set = {e.upper() for e in entities}
        patterns = [(name, pat) for name, pat in _PII_PATTERNS if name in entities_set]
    else:
        patterns = _PII_PATTERNS

    result_text = text
    mapping: dict = {}
    counters: dict = {}

    # Replace from most specific / longest to avoid partial overlaps
    for entity_type, pattern in patterns:
        counter = counters.get(entity_type, 0)

        def _replacer(m, et=entity_type):
            nonlocal counter
            counter += 1
            placeholder = f"[{et}_{counter}]"
            mapping[placeholder] = m.group(0)
            return placeholder

        result_text = pattern.sub(_replacer, result_text)
        counters[entity_type] = counter

    return RedactionResult(
        redacted_text=result_text,
        mapping=mapping,
        entity_counts={k: v for k, v in counters.items() if v > 0},
    )


def rehydrate(text: str, mapping: dict) -> str:
    """
    Restore original PII values from a redaction mapping.

    Replaces placeholders like [EMAIL_1] with their original values.
    Placeholders not found in mapping are left as-is.
    """
    result = text
    # Sort by placeholder length descending to avoid partial replacements
    for placeholder, original in sorted(mapping.items(), key=lambda x: len(x[0]), reverse=True):
        result = result.replace(placeholder, original)
    return result
