"""Privacy accountant utilities for DP unlearning."""

from dataclasses import dataclass, field
from datetime import datetime
from typing import Any


@dataclass
class PrivacySpending:
    """Record of a single privacy-consuming operation."""

    timestamp: datetime
    operation: str
    epsilon: float
    delta: float
    mechanism: str
    metadata: dict[str, Any] = field(default_factory=dict)

    def __repr__(self) -> str:
        return (
            f"PrivacySpending({self.operation}: ε={self.epsilon:.4f}, "
            f"δ={self.delta:.6f}, mechanism={self.mechanism})"
        )


class PrivacyAccountant:
    """Tracks privacy budget consumption across unlearning operations."""

    def __init__(
        self,
        total_epsilon_budget: float,
        total_delta_budget: float = 1e-5,
        composition_type: str = "basic",
    ):
        if total_epsilon_budget <= 0:
            raise ValueError(f"epsilon budget must be positive, got {total_epsilon_budget}")
        if not (0 < total_delta_budget < 1):
            raise ValueError(f"delta budget must be in (0,1), got {total_delta_budget}")

        self.total_epsilon_budget = total_epsilon_budget
        self.total_delta_budget = total_delta_budget
        self.composition_type = composition_type
        self._spending_history: list[PrivacySpending] = []
        self._epsilon_spent = 0.0
        self._delta_spent = 0.0

    def record_operation(
        self,
        epsilon: float,
        delta: float,
        operation: str,
        mechanism: str,
        metadata: dict[str, Any] | None = None,
    ) -> bool:
        new_epsilon, new_delta = self._compute_composition(epsilon, delta)
        if new_epsilon > self.total_epsilon_budget:
            raise ValueError(
                f"Operation would exceed epsilon budget: "
                f"{new_epsilon:.4f} > {self.total_epsilon_budget:.4f}"
            )
        if new_delta > self.total_delta_budget:
            raise ValueError(
                f"Operation would exceed delta budget: "
                f"{new_delta:.6f} > {self.total_delta_budget:.6f}"
            )

        self._spending_history.append(
            PrivacySpending(
                timestamp=datetime.now(),
                operation=operation,
                epsilon=epsilon,
                delta=delta,
                mechanism=mechanism,
                metadata=metadata or {},
            )
        )
        self._epsilon_spent = new_epsilon
        self._delta_spent = new_delta
        return True

    def _compute_composition(self, new_epsilon: float, new_delta: float) -> tuple[float, float]:
        if self.composition_type in {"basic", "advanced", "moments"}:
            return (self._epsilon_spent + new_epsilon, self._delta_spent + new_delta)
        raise ValueError(f"Unknown composition type: {self.composition_type}")

    def get_remaining_budget(self) -> dict[str, float]:
        return {
            "epsilon_remaining": self.total_epsilon_budget - self._epsilon_spent,
            "delta_remaining": self.total_delta_budget - self._delta_spent,
            "epsilon_spent": self._epsilon_spent,
            "delta_spent": self._delta_spent,
        }

    def can_afford(self, epsilon: float, delta: float) -> bool:
        new_epsilon, new_delta = self._compute_composition(epsilon, delta)
        return new_epsilon <= self.total_epsilon_budget and new_delta <= self.total_delta_budget

    def get_spending_history(self) -> list[PrivacySpending]:
        return self._spending_history.copy()

    def reset(self) -> None:
        self._spending_history.clear()
        self._epsilon_spent = 0.0
        self._delta_spent = 0.0

    def summary(self) -> dict[str, Any]:
        return {
            "composition_type": self.composition_type,
            "num_operations": len(self._spending_history),
            **self.get_remaining_budget(),
        }
