from __future__ import annotations

import numpy as np
import pytest

from sage_libs.sage_privacy.dp import GaussianMechanism, LaplaceMechanism
from sage_libs.sage_privacy.unlearning import FisherForgettingUnlearner, GradientAscentUnlearner


def _result_get(result: object, key: str):
    if isinstance(result, dict):
        return result.get(key)
    return getattr(result, key, None)


class TestIssue34DPBoundaries:
    def test_gaussian_rejects_non_positive_sensitivity(self) -> None:
        with pytest.raises(ValueError, match="sensitivity must be positive"):
            GaussianMechanism(epsilon=1.0, sensitivity=0.0)

        with pytest.raises(ValueError, match="sensitivity must be positive"):
            GaussianMechanism(epsilon=1.0, sensitivity=-1.0)

    def test_laplace_rejects_non_positive_sensitivity(self) -> None:
        with pytest.raises(ValueError, match="sensitivity must be positive"):
            LaplaceMechanism(epsilon=1.0, sensitivity=0.0)

        with pytest.raises(ValueError, match="sensitivity must be positive"):
            LaplaceMechanism(epsilon=1.0, sensitivity=-0.5)


class TestIssue34UnlearningKeyPaths:
    def test_gradient_ascent_accepts_non_tuple_forget_data(self) -> None:
        unlearner = GradientAscentUnlearner()
        model = np.random.randn(5)
        forget_data = np.random.randn(8, 5)

        result = unlearner.unlearn(model, forget_data, num_steps=3)

        assert _result_get(result, "success") is True
        assert _result_get(result, "samples_forgotten") == 8

    def test_fisher_rejects_empty_retain_data(self) -> None:
        unlearner = FisherForgettingUnlearner()
        model = np.random.randn(5)
        X_forget = np.random.randn(6, 5)
        y_forget = np.random.randn(6)
        X_retain = np.empty((0, 5))
        y_retain = np.empty((0,))

        with pytest.raises(ValueError, match="retain_data must be non-empty"):
            unlearner.unlearn(
                model,
                (X_forget, y_forget),
                retain_data=(X_retain, y_retain),
            )
