"""Privacy primitives for cross-organization learning.

Uses differential privacy and k-anonymity to enable anonymous
pattern sharing while preserving user/organization privacy.
"""

from __future__ import annotations

import hashlib
import random
import secrets
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

# Check for differential privacy library availability
try:
    from diffprivlib.accountant import BudgetAccountant
    from diffprivlib.mechanisms import Gaussian, Laplace

    DIFFPRIVLIB_AVAILABLE = True
except ImportError:
    DIFFPRIVLIB_AVAILABLE = False
    Laplace = None
    Gaussian = None
    BudgetAccountant = None

# Check for k-anonymity verification library
try:
    from pycanon import anonymity

    PYCANON_AVAILABLE = True
except ImportError:
    PYCANON_AVAILABLE = False
    anonymity = None


class PrivacyBudgetExhausted(Exception):
    """Raised when the differential privacy budget has been fully spent."""


@dataclass
class PrivacyConfig:
    """Privacy configuration for cross-organization sharing.

    Attributes:
        epsilon: Differential privacy budget (lower = more private).
            Typical values: 0.1 (very private) to 10.0 (less private).
        delta: Probability bound for privacy breach (should be very small).
        k_anonymity: Minimum number of organizations before sharing.
        sensitivity: Query sensitivity for DP mechanisms.
        total_budget: Total privacy budget before stopping queries.
    """

    epsilon: float = 1.0
    delta: float = 1e-5
    k_anonymity: int = 3
    sensitivity: float = 1.0
    total_budget: float = 10.0


class DifferentialPrivacy:
    """Apply differential privacy using Diffprivlib.

    Provides mathematically-proven privacy guarantees for aggregate
    statistics through the Laplace and Gaussian mechanisms.
    """

    def __init__(self, config: Optional[PrivacyConfig] = None):
        """Initialize differential privacy mechanism.

        Args:
            config: Privacy configuration settings
        """
        self.config = config or PrivacyConfig()
        self._initialized = False

        # Manual budget tracking (used when diffprivlib is unavailable)
        self._spent_epsilon = 0.0
        self._spent_delta = 0.0

        if DIFFPRIVLIB_AVAILABLE:
            self._init_mechanisms()
        else:
            self.count_mechanism = None
            self.rate_mechanism = None
            self.accountant = None

    def _init_mechanisms(self) -> None:
        """Initialize DP mechanisms with configured parameters."""
        # Laplace mechanism for unbounded counts
        self.count_mechanism = Laplace(
            epsilon=self.config.epsilon, sensitivity=self.config.sensitivity
        )

        # Gaussian mechanism for bounded rates [0, 1]
        self.rate_mechanism = Gaussian(
            epsilon=self.config.epsilon,
            delta=self.config.delta,
            sensitivity=self.config.sensitivity,
        )

        # Budget accountant to track total privacy loss
        self.accountant = BudgetAccountant(
            epsilon=self.config.total_budget, delta=self.config.delta * 10
        )

        self._initialized = True

    @property
    def available(self) -> bool:
        """Check if differential privacy is available."""
        return DIFFPRIVLIB_AVAILABLE

    def _spend_budget(self, epsilon: float, delta: float) -> bool:
        """Register a privacy spend with the budget accountant.

        Args:
            epsilon: Privacy cost of this query
            delta: Delta cost of this query

        Returns:
            True if budget was successfully spent, False if exhausted
        """
        # Always track manually (used for get_remaining_budget fallback)
        self._spent_epsilon += epsilon
        self._spent_delta += delta

        if self.accountant is None:
            return True

        try:
            self.accountant.spend(epsilon, delta)
            return True
        except Exception:
            # BudgetError or other — budget is exhausted
            return False

    def add_noise_to_count(self, true_count: int) -> int:
        """Add Laplacian noise to a count with privacy guarantee.

        Args:
            true_count: The true count value

        Returns:
            Count with differential privacy noise added

        Raises:
            PrivacyBudgetExhausted: If the privacy budget has been fully spent
        """
        # Check and spend budget before applying noise
        if not self.check_budget():
            raise PrivacyBudgetExhausted(
                "Privacy budget exhausted — cannot answer more queries"
            )
        self._spend_budget(self.config.epsilon, self.config.delta)

        if not DIFFPRIVLIB_AVAILABLE or self.count_mechanism is None:
            # Fallback to simple Laplacian noise
            scale = self.config.sensitivity / self.config.epsilon
            noise = self._laplace_sample(scale)
            return max(0, int(true_count + noise))

        try:
            noisy = self.count_mechanism.randomise(float(true_count))
            return max(0, int(round(noisy)))
        except Exception:
            # Fallback on error
            scale = self.config.sensitivity / self.config.epsilon
            noise = self._laplace_sample(scale)
            return max(0, int(true_count + noise))

    def add_noise_to_rate(self, true_rate: float, sample_size: int) -> float:
        """Add noise to a rate while keeping it bounded [0, 1].

        Args:
            true_rate: The true rate value
            sample_size: Number of samples used to compute the rate

        Returns:
            Rate with differential privacy noise, bounded [0, 1]

        Raises:
            PrivacyBudgetExhausted: If the privacy budget has been fully spent
        """
        # Check and spend budget before applying noise
        if not self.check_budget():
            raise PrivacyBudgetExhausted(
                "Privacy budget exhausted — cannot answer more queries"
            )
        self._spend_budget(self.config.epsilon, self.config.delta)

        if not DIFFPRIVLIB_AVAILABLE:
            # Fallback to Gaussian noise (using crypto-safe PRNG for DP guarantees)
            _rng = secrets.SystemRandom()
            scale = self.config.sensitivity / (self.config.epsilon * max(sample_size, 10))
            noise = _rng.gauss(0, scale)
            return max(0.0, min(1.0, true_rate + noise))

        try:
            # Adjust sensitivity based on sample size
            adjusted_sensitivity = self.config.sensitivity / max(sample_size, 10)
            mechanism = Gaussian(
                epsilon=self.config.epsilon,
                delta=self.config.delta,
                sensitivity=adjusted_sensitivity,
            )
            noisy = mechanism.randomise(true_rate)
            return max(0.0, min(1.0, noisy))
        except Exception:
            # Fallback on error (using crypto-safe PRNG for DP guarantees)
            _rng = secrets.SystemRandom()
            scale = self.config.sensitivity / (self.config.epsilon * max(sample_size, 10))
            noise = _rng.gauss(0, scale)
            return max(0.0, min(1.0, true_rate + noise))

    def _laplace_sample(self, scale: float) -> float:
        """Sample from Laplace distribution.

        Args:
            scale: Scale parameter (b) of the Laplace distribution

        Returns:
            Random sample from Laplace(0, scale)
        """
        import math

        u = secrets.SystemRandom().random() - 0.5
        if u == 0:
            return 0.0
        sign = 1 if u >= 0 else -1
        return -scale * sign * math.log(1 - 2 * abs(u) + 1e-10)

    def get_remaining_budget(self) -> float:
        """Get remaining privacy budget.

        Returns:
            Remaining epsilon budget
        """
        if self.accountant:
            try:
                spent = self.accountant.total()
                return max(0, self.config.total_budget - spent[0])
            except Exception:
                pass
        # Fallback: use manual tracking
        return max(0.0, self.config.total_budget - self._spent_epsilon)

    def check_budget(self) -> bool:
        """Check if there is remaining privacy budget.

        Returns:
            True if queries can still be made
        """
        remaining = self.get_remaining_budget()
        return remaining > self.config.epsilon


class KAnonymityChecker:
    """Verify k-anonymity for shared data.

    Ensures that each pattern group has at least k organizations
    contributing before it can be shared.
    """

    def __init__(self, k: int = 3):
        """Initialize k-anonymity checker.

        Args:
            k: Minimum group size for anonymity (default: 3)
        """
        self.k = k

    @property
    def available(self) -> bool:
        """Check if pyCANON is available."""
        return PYCANON_AVAILABLE

    def is_k_anonymous(
        self, records: List[Dict[str, Any]], quasi_identifiers: List[str]
    ) -> bool:
        """Check if records satisfy k-anonymity.

        Args:
            records: List of records to check
            quasi_identifiers: Fields that together might identify a record

        Returns:
            True if all equivalence classes have at least k records
        """
        if not records:
            return False

        if PYCANON_AVAILABLE and anonymity is not None:
            return self._pycanon_k_check(records, quasi_identifiers)

        return self._simple_k_check(records, quasi_identifiers)

    def _pycanon_k_check(
        self, records: List[Dict[str, Any]], quasi_identifiers: List[str]
    ) -> bool:
        """Check k-anonymity using pyCANON library.

        Args:
            records: List of records
            quasi_identifiers: Quasi-identifier fields

        Returns:
            True if k-anonymous
        """
        try:
            import pandas as pd

            df = pd.DataFrame(records)
            # Only check columns that exist
            valid_qis = [qi for qi in quasi_identifiers if qi in df.columns]
            if not valid_qis:
                return False
            return anonymity.k_anonymity(df, valid_qis) >= self.k
        except Exception:
            return self._simple_k_check(records, quasi_identifiers)

    def _simple_k_check(
        self, records: List[Dict[str, Any]], quasi_identifiers: List[str]
    ) -> bool:
        """Simple k-anonymity check without pyCANON.

        Args:
            records: List of records
            quasi_identifiers: Quasi-identifier fields

        Returns:
            True if all equivalence classes have >= k records
        """
        from collections import Counter

        # Group by quasi-identifiers
        groups: Counter = Counter()
        for record in records:
            key = tuple(str(record.get(qi, "")) for qi in quasi_identifiers)
            groups[key] += 1

        # Check all groups have at least k records
        if not groups:
            return False

        return all(count >= self.k for count in groups.values())

    def get_equivalence_class_sizes(
        self, records: List[Dict[str, Any]], quasi_identifiers: List[str]
    ) -> Dict[str, int]:
        """Get size of each equivalence class.

        Args:
            records: List of records
            quasi_identifiers: Quasi-identifier fields

        Returns:
            Dict mapping class key to count
        """
        from collections import Counter

        groups: Counter = Counter()
        for record in records:
            key = str(tuple(str(record.get(qi, "")) for qi in quasi_identifiers))
            groups[key] += 1

        return dict(groups)

    def find_small_classes(
        self, records: List[Dict[str, Any]], quasi_identifiers: List[str]
    ) -> List[str]:
        """Find equivalence classes smaller than k.

        Args:
            records: List of records
            quasi_identifiers: Quasi-identifier fields

        Returns:
            List of class keys with fewer than k records
        """
        sizes = self.get_equivalence_class_sizes(records, quasi_identifiers)
        return [key for key, count in sizes.items() if count < self.k]


def hash_identifier(identifier: str, salt: str = "") -> str:
    """Hash an identifier for privacy using HMAC.

    Args:
        identifier: The identifier to hash
        salt: Optional salt/key for the HMAC

    Returns:
        32-character hex hash (128-bit collision resistance)
    """
    import hmac
    return hmac.new(
        salt.encode(), identifier.encode(), hashlib.sha256
    ).hexdigest()[:32]


def hash_code_pattern(code: str) -> str:
    """Hash code to a structural pattern hash.

    Args:
        code: Code string to hash

    Returns:
        16-character hex hash of normalized code structure
    """
    # Normalize whitespace and common patterns
    import re

    normalized = code.strip()
    normalized = re.sub(r"\s+", " ", normalized)
    normalized = re.sub(r"['\"].*?['\"]", "STR", normalized)
    normalized = re.sub(r"\b\d+\b", "NUM", normalized)

    return hashlib.sha256(normalized.encode()).hexdigest()[:16]


@dataclass
class PrivacyMetrics:
    """Metrics about privacy preservation."""

    total_patterns_shared: int = 0
    patterns_suppressed: int = 0  # Below k-anonymity threshold
    avg_equivalence_class_size: float = 0.0
    noise_added_count: int = 0
    privacy_budget_remaining: float = 0.0


class PrivacyAuditor:
    """Audit privacy of shared data.

    Tracks what data is being shared and ensures privacy
    guarantees are maintained.
    """

    def __init__(self, config: Optional[PrivacyConfig] = None):
        """Initialize privacy auditor.

        Args:
            config: Privacy configuration
        """
        self.config = config or PrivacyConfig()
        self.dp = DifferentialPrivacy(config)
        self.k_checker = KAnonymityChecker(self.config.k_anonymity)
        self.metrics = PrivacyMetrics()

    def audit_pattern(self, pattern_data: Dict[str, Any]) -> Dict[str, Any]:
        """Audit a pattern before sharing.

        Args:
            pattern_data: Pattern data to audit

        Returns:
            Dict with audit results
        """
        issues = []

        # Check for potentially identifying information
        if self._contains_paths(pattern_data):
            issues.append("Contains file paths")

        if self._contains_user_data(pattern_data):
            issues.append("Contains potential user data")

        if self._contains_credentials(pattern_data):
            issues.append("Contains potential credentials")

        return {
            "safe_to_share": len(issues) == 0,
            "issues": issues,
            "privacy_score": 1.0 - (len(issues) * 0.25),
        }

    def _contains_paths(self, data: Dict[str, Any]) -> bool:
        """Check if data contains file paths."""
        import re

        text = str(data)
        path_pattern = r"(/[^\s]+/[^\s]+|[A-Z]:\\[^\s]+)"
        return bool(re.search(path_pattern, text))

    def _contains_user_data(self, data: Dict[str, Any]) -> bool:
        """Check if data might contain user information."""
        import re

        text = str(data).lower()
        patterns = [
            r"user[_\s]?id",
            r"email",
            r"@[a-z]+\.[a-z]+",  # Email pattern
            r"user[_\s]?name",
        ]
        return any(re.search(p, text) for p in patterns)

    def _contains_credentials(self, data: Dict[str, Any]) -> bool:
        """Check if data might contain credentials."""
        text = str(data).lower()
        keywords = ["password", "secret", "api_key", "token", "credential", "auth"]
        return any(kw in text for kw in keywords)

    def get_metrics(self) -> PrivacyMetrics:
        """Get current privacy metrics.

        Returns:
            PrivacyMetrics with current stats
        """
        self.metrics.privacy_budget_remaining = self.dp.get_remaining_budget()
        return self.metrics
