"""Tests for v0.37.0 features."""
import math
from unittest.mock import MagicMock, patch
import pytest
from llm_guard.direct_generation_guard import DirectGenerationGuard, DirectGenResult


class TestSelfConsistency:
    def _make_guard(self):
        g = DirectGenerationGuard(api_key="fake-key")
        return g

    def test_jaccard_identical(self):
        g = self._make_guard()
        assert g._jaccard("the cat sat", "the cat sat") == 1.0

    def test_jaccard_disjoint(self):
        g = self._make_guard()
        assert g._jaccard("hello world", "foo bar") == 0.0

    def test_jaccard_partial(self):
        g = self._make_guard()
        j = g._jaccard("the cat", "the dog")
        assert 0.0 < j < 1.0

    def test_sc_risk_high_for_diverse_answers(self):
        """When sampled answers are all different, sc_risk should be high."""
        g = self._make_guard()

        base_result = DirectGenResult(
            risk_score=0.4,
            confidence_tier="HIGH",
            ptrue_score=0.6,
            lexical_risk=0.3,
            signals={},
        )
        sampled = [
            "Paris is a city in France.",
            "The capital of France is Paris.",
            "London is the capital of England.",
            "Berlin is a major European city.",
            "Rome was founded by Romulus.",
        ]

        with patch.object(g, "score", return_value=base_result):
            with patch.object(g, "_generate_answer", side_effect=sampled):
                result = g.score_with_selfconsistency("What is Paris?", "Paris is in France.")

        assert "sc_risk" in result.signals
        assert "sc_mean_sim" in result.signals
        assert result.signals["sc_risk"] > 0.5  # diverse answers → high sc_risk
        assert 0.0 <= result.risk_score <= 1.0

    def test_sc_risk_low_for_consistent_answers(self):
        """When all answers are identical, sc_risk should be near 0."""
        g = self._make_guard()

        base_result = DirectGenResult(
            risk_score=0.2,
            confidence_tier="HIGH",
            ptrue_score=0.8,
            lexical_risk=0.2,
            signals={},
        )
        identical = ["Paris is the capital of France."] * 5

        with patch.object(g, "score", return_value=base_result):
            with patch.object(g, "_generate_answer", side_effect=identical):
                result = g.score_with_selfconsistency("What is the capital of France?", "Paris is the capital.")

        assert result.signals["sc_risk"] < 0.2  # near 0 for identical answers

    def test_score_with_selfconsistency_no_api_key(self):
        """Without API key, should raise RuntimeError."""
        g = DirectGenerationGuard(api_key="")
        with pytest.raises(RuntimeError, match="api_key"):
            g.score_with_selfconsistency("q", "a")

    def test_result_has_sc_risk_field(self):
        """DirectGenResult should have sc_risk field."""
        r = DirectGenResult(risk_score=0.5, confidence_tier="MEDIUM",
                            ptrue_score=0.5, lexical_risk=0.5, signals={})
        assert hasattr(r, "sc_risk")
        assert math.isnan(r.sc_risk)

    def test_n_samples_zero_raises(self):
        """n_samples=0 should raise ValueError."""
        g = self._make_guard()
        with pytest.raises(ValueError, match="n_samples"):
            g.score_with_selfconsistency("q", "a", n_samples=0)

    def test_n_samples_negative_raises(self):
        """n_samples<0 should raise ValueError."""
        g = self._make_guard()
        with pytest.raises(ValueError, match="n_samples"):
            g.score_with_selfconsistency("q", "a", n_samples=-1)

    def test_n_samples_one_produces_single_pair(self):
        """n_samples=1 should work — exactly 1 pair, result is valid."""
        g = self._make_guard()
        base_result = DirectGenResult(
            risk_score=0.3, confidence_tier="HIGH",
            ptrue_score=0.7, lexical_risk=0.3, signals={},
        )
        with patch.object(g, "score", return_value=base_result):
            with patch.object(g, "_generate_answer", return_value="Some answer."):
                result = g.score_with_selfconsistency("q", "Some answer.", n_samples=1)
        assert "sc_risk" in result.signals
        assert 0.0 <= result.risk_score <= 1.0


class TestQuickRecalibrate:
    def _make_guard(self):
        from llm_guard.agent_guard import AgentGuard
        return AgentGuard.__new__(AgentGuard)

    def test_quick_recalibrate_insufficient_chains_raises(self):
        from llm_guard.agent_guard import AgentGuard
        g = AgentGuard()
        chains = [{"question": "q", "steps": [], "final_answer": "a"}] * 5
        labels = [0] * 5
        with pytest.raises(ValueError, match="min_chains"):
            g.quick_recalibrate(chains, labels, min_chains=20)

    def test_quick_recalibrate_mismatched_lengths_raises(self):
        from llm_guard.agent_guard import AgentGuard
        g = AgentGuard()
        chains = [{"question": "q", "steps": [], "final_answer": "a"}] * 5
        labels = [0] * 4
        with pytest.raises(ValueError, match="same length"):
            g.quick_recalibrate(chains, labels)

    def test_quick_recalibrate_fits_calibrator(self):
        """quick_recalibrate with 20+ chains should set _iso_calibrator."""
        from unittest.mock import patch, MagicMock
        from llm_guard.agent_guard import AgentGuard
        g = AgentGuard()
        chains = [{"question": "q", "steps": [], "final_answer": "a"}] * 20
        labels = [i % 2 for i in range(20)]  # alternating 0/1

        mock_result = MagicMock()
        mock_result.risk_score = 0.5

        with patch.object(g, "score_chain", return_value=mock_result):
            result = g.quick_recalibrate(chains, labels)

        assert result is g  # method chaining
        assert g._iso_calibrator is not None

    def test_check_calibration_staleness_no_calibrator(self):
        from llm_guard.agent_guard import AgentGuard
        g = AgentGuard()
        assert g.check_calibration_staleness() is False

    def test_check_calibration_staleness_stale(self):
        """Changing judge_model after calibration should trigger a warning."""
        import warnings
        from unittest.mock import MagicMock
        from llm_guard.agent_guard import AgentGuard
        g = AgentGuard()
        g._iso_calibrator = MagicMock()
        g._calibrated_judge_model = "model-a"
        g._judge_model = "model-b"
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            stale = g.check_calibration_staleness()
        assert stale is True
        assert len(w) == 1
        assert "model-a" in str(w[0].message)

    def test_check_calibration_staleness_not_stale(self):
        from unittest.mock import MagicMock
        from llm_guard.agent_guard import AgentGuard
        g = AgentGuard()
        g._iso_calibrator = MagicMock()
        g._calibrated_judge_model = "model-a"
        g._judge_model = "model-a"
        assert g.check_calibration_staleness() is False

    def test_calibrated_judge_model_set_on_calibrate_isotonic(self):
        """calibrate_isotonic() should record the judge model."""
        from llm_guard.agent_guard import AgentGuard
        g = AgentGuard()
        g.calibrate_isotonic([0.2, 0.5, 0.8], [0, 1, 1])
        assert g._calibrated_judge_model == g._judge_model


class TestAsyncMethods:
    def test_ascore_chain_is_coroutine(self):
        import asyncio
        import inspect
        from llm_guard.agent_guard import AgentGuard
        g = AgentGuard()
        coro = g.ascore_chain("q", [], "a")
        assert inspect.iscoroutine(coro)
        coro.close()  # clean up without running

    def test_ascore_chain_runs(self):
        """ascore_chain should return a ChainTrustResult."""
        import asyncio
        from unittest.mock import patch, MagicMock
        from llm_guard.agent_guard import AgentGuard, ChainTrustResult

        g = AgentGuard()
        mock = MagicMock(spec=ChainTrustResult)

        async def run():
            with patch.object(g, "score_chain", return_value=mock) as m:
                result = await g.ascore_chain("q", [], "a")
                m.assert_called_once_with("q", [], "a")
                return result

        result = asyncio.run(run())
        assert result is mock

    def test_atrack_is_coroutine(self):
        import asyncio
        import inspect
        from llm_guard.agent_guard import AgentGuard
        g = AgentGuard()
        coro = g.atrack("q", [], "a")
        assert inspect.iscoroutine(coro)
        coro.close()
