"""End-to-end DP unlearning engine."""

from dataclasses import dataclass, field
from typing import Any

import numpy as np

from .base_mechanism import BasePrivacyMechanism, SimpleLaplaceMechanism
from .neighbor_compensation import NeighborCompensation
from .privacy_accountant import PrivacyAccountant
from .vector_perturbation import VectorPerturbation


@dataclass
class UnlearningResult:
    """Result of an unlearning operation."""

    success: bool
    num_vectors_unlearned: int
    num_neighbors_compensated: int
    privacy_cost: tuple[float, float]
    metadata: dict[str, Any] = field(default_factory=dict)


class UnlearningEngine:
    """Main engine for differential privacy-based machine unlearning."""

    def __init__(
        self,
        epsilon: float = 1.0,
        delta: float = 1e-5,
        total_budget_epsilon: float = 10.0,
        total_budget_delta: float = 1e-4,
        mechanism: BasePrivacyMechanism | None = None,
        enable_compensation: bool = True,
    ):
        self.mechanism = mechanism or SimpleLaplaceMechanism(epsilon=epsilon)
        self.privacy_accountant = PrivacyAccountant(
            total_epsilon_budget=total_budget_epsilon,
            total_delta_budget=total_budget_delta,
        )
        self.vector_perturbation = VectorPerturbation(self.mechanism)
        self.neighbor_compensation = NeighborCompensation() if enable_compensation else None
        self.enable_compensation = enable_compensation
        self.epsilon = epsilon
        self.delta = delta

    def unlearn_vectors(
        self,
        vectors_to_forget: np.ndarray,
        vector_ids_to_forget: list[str],
        all_vectors: np.ndarray | None = None,
        all_vector_ids: list[str] | None = None,
        perturbation_strategy: str = "uniform",
        return_compensated_neighbors: bool = False,
    ) -> UnlearningResult:
        n_forget = len(vectors_to_forget)
        operation_epsilon = self.epsilon * n_forget
        operation_delta = self.delta * n_forget

        if not self.privacy_accountant.can_afford(operation_epsilon, operation_delta):
            return UnlearningResult(
                success=False,
                num_vectors_unlearned=0,
                num_neighbors_compensated=0,
                privacy_cost=(0.0, 0.0),
                metadata={
                    "error": "Insufficient privacy budget",
                    "remaining_budget": self.privacy_accountant.get_remaining_budget(),
                },
            )

        perturbed_vectors = self.vector_perturbation.perturb_batch_vectors(
            vectors_to_forget,
            strategy=perturbation_strategy,
        )

        compensated_neighbors: dict[str, np.ndarray] = {}
        num_compensated = 0
        if (
            self.enable_compensation
            and self.neighbor_compensation is not None
            and all_vectors is not None
            and all_vector_ids is not None
        ):
            for original, perturbed, _vec_id in zip(
                vectors_to_forget,
                perturbed_vectors,
                vector_ids_to_forget,
                strict=False,
            ):
                neighbor_compensations = self.neighbor_compensation.apply_compensation(
                    original,
                    perturbed,
                    all_vectors,
                    all_vector_ids,
                )
                compensated_neighbors.update(neighbor_compensations)
                num_compensated += len(neighbor_compensations)

        self.privacy_accountant.record_operation(
            epsilon=operation_epsilon,
            delta=operation_delta,
            operation=f"unlearn_{n_forget}_vectors",
            mechanism=self.mechanism.name,
            metadata={
                "num_vectors": n_forget,
                "perturbation_strategy": perturbation_strategy,
                "compensation_enabled": self.enable_compensation,
                "num_compensated": num_compensated,
            },
        )

        result = UnlearningResult(
            success=True,
            num_vectors_unlearned=n_forget,
            num_neighbors_compensated=num_compensated,
            privacy_cost=(operation_epsilon, operation_delta),
            metadata={
                "perturbation_strategy": perturbation_strategy,
                "perturbed_vectors": perturbed_vectors,
                "privacy_accountant_summary": self.privacy_accountant.summary(),
            },
        )
        if return_compensated_neighbors:
            result.metadata["compensated_neighbors"] = compensated_neighbors
        return result
