"""tests/test_calibration_conformal_risk_budget.py — Unit tests for llm_guard.calibration.conformal_risk_budget."""
import pytest

pytestmark = pytest.mark.unit


class TestImport:
    def test_import(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        assert ConformalRiskBudget is not None

    def test_import_budget_state(self):
        from llm_guard.calibration.conformal_risk_budget import BudgetState
        assert BudgetState is not None


class TestInstantiation:
    def test_default_construction(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        b = ConformalRiskBudget()
        assert b._budget == 0.20

    def test_bonferroni_method(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        b = ConformalRiskBudget(method="bonferroni")
        assert b._method == "bonferroni"

    def test_product_method(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        b = ConformalRiskBudget(method="product")
        assert b._method == "product"

    def test_invalid_budget_raises(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        with pytest.raises(ValueError):
            ConformalRiskBudget(total_budget=0.0)
        with pytest.raises(ValueError):
            ConformalRiskBudget(total_budget=1.0)


class TestAllocate:
    def test_allocate_bonferroni(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        b = ConformalRiskBudget(total_budget=0.20, method="bonferroni")
        thresholds = b.allocate(5)
        assert len(thresholds) == 5
        assert all(abs(t - 0.04) < 1e-9 for t in thresholds)

    def test_allocate_product(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        b = ConformalRiskBudget(total_budget=0.20, method="product")
        thresholds = b.allocate(5)
        assert len(thresholds) == 5
        assert all(t > 0 for t in thresholds)

    def test_allocate_invalid_steps(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        b = ConformalRiskBudget()
        with pytest.raises(ValueError):
            b.allocate(0)


class TestObserve:
    def test_observe_before_allocate_raises(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        b = ConformalRiskBudget()
        with pytest.raises(RuntimeError):
            b.observe(0, 0.1)

    def test_observe_returns_false_when_under_budget(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        b = ConformalRiskBudget(total_budget=0.20, method="bonferroni")
        b.allocate(5)
        abort = b.observe(0, 0.01)
        assert abort is False

    def test_observe_returns_true_when_over_budget(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        b = ConformalRiskBudget(total_budget=0.20, method="bonferroni")
        b.allocate(5)
        b.observe(0, 0.15)
        abort = b.observe(1, 0.15)
        assert abort is True

    def test_remaining_budget(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        b = ConformalRiskBudget(total_budget=0.20)
        b.allocate(5)
        b.observe(0, 0.05)
        assert abs(b.remaining_budget - 0.15) < 1e-9

    def test_summary_keys(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        b = ConformalRiskBudget(total_budget=0.20)
        b.allocate(3)
        b.observe(0, 0.05)
        s = b.summary()
        assert "total_budget" in s
        assert "n_observed" in s
        assert "over_budget" in s

    def test_reset(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        b = ConformalRiskBudget(total_budget=0.20)
        b.allocate(5)
        b.observe(0, 0.15)
        b.reset()
        assert b.remaining_budget == 0.20


class TestAdaptiveMethod:
    def test_adaptive_recomputes_threshold(self):
        from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget
        b = ConformalRiskBudget(total_budget=0.20, method="adaptive")
        b.allocate(5)
        abort = b.observe(0, 0.05)
        assert isinstance(abort, bool)
