"""
FoundryMatch SaaS - Enhanced Privacy Filter
===========================================
Provides a robust, pattern-based PII (Personally Identifiable Information)
filter to help ensure data privacy and compliance.
"""

import hashlib
import logging
import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Literal, Optional

import pandas as pd

log = logging.getLogger(__name__)

PIIAction = Literal["mask", "hash", "redact"]


@dataclass
class PIIPattern:
    """Defines a pattern for detecting a specific type of PII."""

    name: str
    regex: re.Pattern
    action: PIIAction = "mask"
    context_keywords: List[str] = field(default_factory=list)


class EnhancedPrivacyFilter:
    """
    Scans and sanitizes data (DataFrames or dicts) to remove or anonymize PII.
    """

    def __init__(self, salt: str, custom_patterns: Optional[List[PIIPattern]] = None):
        """
        Initializes the filter with a salt for hashing and optional custom patterns.

        Args:
            salt: A secret salt to be used for hashing PII.
            custom_patterns: A list of PIIPattern objects to extend the default set.
        """
        if not salt:
            raise ValueError("A salt is required for hashing PII.")
        self.salt = salt.encode("utf-8")
        self.patterns = self._load_default_patterns()
        if custom_patterns:
            self.patterns.extend(custom_patterns)

    def _load_default_patterns(self) -> List[PIIPattern]:
        """Loads a default set of common PII patterns."""
        return [
            PIIPattern(
                name="EMAIL",
                regex=re.compile(r"[\w\.\-]{1,256}@[\w\.\-]{1,256}\.[\w]{2,4}"),
                action="mask",
                context_keywords=["email", "mail", "user_id"],
            ),
            PIIPattern(
                name="PHONE",
                regex=re.compile(r"\(?\d{3}\)?[\s\-\.]?\d{3}[\s\-\.]?\d{4}"),
                action="mask",
                context_keywords=["phone", "mobile", "fax", "contact"],
            ),
            PIIPattern(
                name="CREDIT_CARD",
                regex=re.compile(r"\b(?:\d[ -]*?){13,16}\b"),
                action="redact",
                context_keywords=["card", "cc", "payment", "credit", "account"],
            ),
            PIIPattern(
                name="IP_ADDRESS",
                regex=re.compile(r"\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b"),
                action="hash",
                context_keywords=["ip", "address", "network"],
            ),
            PIIPattern(
                name="SSN",
                regex=re.compile(r"\b\d{3}-\d{2}-\d{4}\b"),
                action="redact",
                context_keywords=["ssn", "social_security"],
            ),
        ]

    def _transform_value(self, value: str, pattern: PIIPattern) -> str:
        """Applies the specified transformation action to a value."""
        action = pattern.action
        if action == "redact":
            return f"[{pattern.name}_REDACTED]"

        if action == "hash":
            hashed = hashlib.sha256(self.salt + str(value).encode("utf-8")).hexdigest()
            return f"hash:{pattern.name}:{hashed[:16]}"

        if action == "mask":
            if pattern.name == "EMAIL":
                parts = str(value).split("@")
                if len(parts) == 2:
                    user, domain = parts
                    return f"{user[0]}***@{domain[0]}***.{domain.split('.')[-1]}"
            elif pattern.name == "PHONE":
                digits = re.sub(r"\D", "", str(value))
                if len(digits) >= 10:
                    return f"***-***-{digits[-4:]}"
            # Default mask
            return f"{str(value)[0]}***{str(value)[-1]}"

        return value

    def _scan_and_transform(self, value: Any) -> Any:
        """Scans a single value and transforms it if it matches any PII pattern."""
        if not isinstance(value, str) or not value:
            return value

        for pattern in self.patterns:
            if pattern.regex.search(value):
                # Found a match, transform and return immediately
                return self._transform_value(value, pattern)

        return value

    def filter_dataframe(
        self, df: pd.DataFrame, sample_size: int = 100
    ) -> pd.DataFrame:
        """
        Scans and filters a pandas DataFrame for PII.

        It first checks column names for context keywords and then scans a sample
        of the data to find columns containing PII.

        Args:
            df: The DataFrame to filter.
            sample_size: The number of rows to sample for pattern matching.

        Returns:
            A new DataFrame with PII transformed.
        """
        df_copy = df.copy()
        sample = df_copy.head(sample_size)

        for col in df_copy.columns:
            if df_copy[col].dtype != "object":
                continue

            col_lower = col.lower()
            patterns_to_check = [
                p
                for p in self.patterns
                if any(kw in col_lower for kw in p.context_keywords)
            ]

            # If no context match, check all patterns on a sample
            if not patterns_to_check:
                if any(
                    sample[col].astype(str).str.contains(p.regex, na=False).any()
                    for p in self.patterns
                ):
                    patterns_to_check = self.patterns

            if not patterns_to_check:
                continue

            log.info(f"Applying PII filter to column '{col}'...")

            # Apply all relevant patterns to the column
            def apply_all_patterns(val):
                if not isinstance(val, str):
                    return val
                for p in patterns_to_check:
                    if p.regex.search(val):
                        return self._transform_value(val, p)
                return val

            df_copy[col] = df_copy[col].apply(apply_all_patterns)

        return df_copy

    def filter_payload(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        """
        Recursively scans and filters a dictionary payload for PII.

        Args:
            payload: The dictionary to filter.

        Returns:
            A new dictionary with PII transformed.
        """
        filtered_payload = {}
        for key, value in payload.items():
            if isinstance(value, dict):
                filtered_payload[key] = self.filter_payload(value)
            elif isinstance(value, list):
                filtered_payload[key] = [
                    self.filter_payload(item)
                    if isinstance(item, dict)
                    else self._scan_and_transform(item)
                    for item in value
                ]
            else:
                filtered_payload[key] = self._scan_and_transform(value)
        return filtered_payload
