"""tests/test_calibration_zero_shot.py — Unit tests for llm_guard.calibration.zero_shot_calibrator."""
import pytest
from unittest.mock import MagicMock

pytestmark = pytest.mark.unit


def _make_chains(n=10):
    return [{"question": f"Q{i}", "steps": [], "final_answer": f"A{i}"} for i in range(n)]


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


class TestInstantiation:
    def test_construction_with_mock_guard(self):
        from llm_guard.calibration.zero_shot_calibrator import ZeroShotCalibrator
        mock_guard = MagicMock()
        cal = ZeroShotCalibrator(guard=mock_guard)
        assert not cal.is_fitted

    def test_ptrue_threshold(self):
        from llm_guard.calibration.zero_shot_calibrator import ZeroShotCalibrator
        cal = ZeroShotCalibrator(guard=MagicMock(), ptrue_threshold=0.6)
        assert cal.ptrue_threshold == 0.6

    def test_n_pseudo_zero_initially(self):
        from llm_guard.calibration.zero_shot_calibrator import ZeroShotCalibrator
        cal = ZeroShotCalibrator(guard=MagicMock())
        assert cal.n_pseudo == 0


class TestFitValidation:
    def test_empty_chains_raises(self):
        from llm_guard.calibration.zero_shot_calibrator import ZeroShotCalibrator
        cal = ZeroShotCalibrator(guard=MagicMock())
        with pytest.raises(ValueError, match="empty"):
            cal.fit([], n_pseudo=10)

    def test_too_few_pseudo_raises(self):
        from llm_guard.calibration.zero_shot_calibrator import ZeroShotCalibrator
        cal = ZeroShotCalibrator(guard=MagicMock())
        with pytest.raises(ValueError, match="n_pseudo"):
            cal.fit(_make_chains(10), n_pseudo=3)


class TestFitWithMock:
    def _make_fitted_cal(self, n=10):
        from llm_guard.calibration.zero_shot_calibrator import ZeroShotCalibrator

        mock_guard = MagicMock()
        # Mock score_with_ptrue to return a result with behavioral_score and behavioral_components
        mock_result = MagicMock()
        mock_result.behavioral_score = 0.4
        mock_result.behavioral_components = {"ptrue_risk": 0.3}
        mock_guard.score_with_ptrue.return_value = mock_result
        mock_guard.calibrate_isotonic = MagicMock()

        cal = ZeroShotCalibrator(guard=mock_guard)
        chains = _make_chains(n)
        cal.fit(chains, n_pseudo=n)
        return cal, mock_guard

    def test_fit_marks_fitted(self):
        cal, _ = self._make_fitted_cal(10)
        assert cal.is_fitted

    def test_get_threshold_after_fit(self):
        cal, _ = self._make_fitted_cal(10)
        # All pseudo-labels should be 0 (risk=0.3 < threshold 0.55)
        threshold = cal.get_threshold(target_fpr=0.10)
        assert 0.0 <= threshold <= 1.0

    def test_get_threshold_raises_unfitted(self):
        from llm_guard.calibration.zero_shot_calibrator import ZeroShotCalibrator
        cal = ZeroShotCalibrator(guard=MagicMock())
        with pytest.raises(RuntimeError):
            cal.get_threshold()

    def test_calibration_quality_estimate(self):
        cal, _ = self._make_fitted_cal(10)
        quality = cal.calibration_quality_estimate()
        assert "quality" in quality
        assert quality["quality"] in ("high", "medium", "low")
        assert "n_pseudo" in quality

    def test_repr(self):
        from llm_guard.calibration.zero_shot_calibrator import ZeroShotCalibrator
        cal = ZeroShotCalibrator(guard=MagicMock())
        r = repr(cal)
        assert "ZeroShotCalibrator" in r
