"""tests/test_calibration_silent_failure.py — Unit tests for llm_guard.calibration.silent_failure_calibrator."""
import pytest
from unittest.mock import MagicMock
import numpy as np

pytestmark = pytest.mark.unit


def _make_chains(n=20, frac_wrong=0.5):
    chains = []
    for i in range(n):
        chains.append({
            "question": "Who is X?",
            "steps": [{"thought": "searching", "action_type": "Search",
                        "observation": "X is someone.", "action_arg": "X"}],
            "final_answer": "X is someone.",
            "correct": i >= int(n * frac_wrong),
        })
    return chains


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


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

    def test_default_construction(self):
        from llm_guard.calibration.silent_failure_calibrator import SilentFailureCalibrator
        # Creates a real AgentGuard internally
        cal = SilentFailureCalibrator()
        assert cal._C == 0.1

    def test_feat_keys_length(self):
        from llm_guard.calibration.silent_failure_calibrator import SilentFailureCalibrator
        assert len(SilentFailureCalibrator.FEAT_KEYS) == 12


class TestFit:
    def _make_mock_guard(self):
        mock_guard = MagicMock()
        mock_result = MagicMock()
        # Return a behavioral_components dict with all required features
        mock_result.behavioral_components = {
            k: 0.5 for k in [
                "sc1", "sc2", "sc3", "sc5", "sc6", "sc8",
                "sc9", "sc10", "sc11", "sc12", "mean_sim", "search_count_norm"
            ]
        }
        mock_guard.score_chain.return_value = mock_result
        return mock_guard

    def test_fit_too_few_labels_raises(self):
        from llm_guard.calibration.silent_failure_calibrator import SilentFailureCalibrator
        cal = SilentFailureCalibrator(guard=self._make_mock_guard())
        # Only wrong chains → only one class
        chains = [c for c in _make_chains(15) if not c["correct"]]
        with pytest.raises(ValueError, match="both correct and wrong"):
            cal.fit(chains)

    def test_fit_marks_fitted(self):
        from llm_guard.calibration.silent_failure_calibrator import SilentFailureCalibrator
        cal = SilentFailureCalibrator(guard=self._make_mock_guard())
        cal.fit(_make_chains(20))
        assert cal.is_fitted

    def test_score_after_fit(self):
        from llm_guard.calibration.silent_failure_calibrator import SilentFailureCalibrator
        cal = SilentFailureCalibrator(guard=self._make_mock_guard())
        cal.fit(_make_chains(20))
        risk = cal.score("Q", [], "A")
        assert 0.0 <= risk <= 1.0

    def test_score_raises_unfitted(self):
        from llm_guard.calibration.silent_failure_calibrator import SilentFailureCalibrator
        cal = SilentFailureCalibrator(guard=self._make_mock_guard())
        with pytest.raises(RuntimeError):
            cal.score("Q", [], "A")

    def test_feature_importances(self):
        from llm_guard.calibration.silent_failure_calibrator import SilentFailureCalibrator
        cal = SilentFailureCalibrator(guard=self._make_mock_guard())
        cal.fit(_make_chains(20))
        fi = cal.feature_importances()
        assert len(fi) == 12

    def test_n_train(self):
        from llm_guard.calibration.silent_failure_calibrator import SilentFailureCalibrator
        cal = SilentFailureCalibrator(guard=self._make_mock_guard())
        cal.fit(_make_chains(20))
        assert cal.n_train == 20

    def test_repr(self):
        from llm_guard.calibration.silent_failure_calibrator import SilentFailureCalibrator
        cal = SilentFailureCalibrator(guard=self._make_mock_guard())
        r = repr(cal)
        assert "SilentFailureCalibrator" in r
