# tests/test_chat_guard_v027.py
"""Tests for v0.27.0 ChatGuard additions: auto_fit(), fit(cov_features_list=...), 16-feature path."""
import pytest
import logging
from unittest.mock import patch, MagicMock
from llm_guard.chat_guard import ChatGuard, ChatResult, COV_FEATURE_KEYS
from llm_guard.chat_cov_features import CoVFeatureVector

# ── Fixtures ──────────────────────────────────────────────────────────────────

LABELED_PAIRS_20 = (
    [{"query": f"q{i}", "response": "Paris is the capital.", "label": 0} for i in range(10)]
    + [{"query": f"q{i}", "response": "I think maybe not sure.", "label": 1} for i in range(10)]
)

COV_DICT = CoVFeatureVector(0.1, 0.5, 0.2, 0.05).to_dict()
COV_FEATURES_20 = [COV_DICT] * 20


# ── COV_FEATURE_KEYS constant ─────────────────────────────────────────────────

class TestCOVFeatureKeys:
    def test_cov_feature_keys_has_four_elements(self):
        assert len(COV_FEATURE_KEYS) == 4

    def test_cov_feature_keys_match_cov_vector_dict(self):
        vec = CoVFeatureVector(0.1, 0.5, 0.2, 0.05)
        assert set(COV_FEATURE_KEYS) == set(vec.to_dict().keys())


# ── fit() with cov_features_list ─────────────────────────────────────────────

class TestFitWithCoVFeatures:
    def test_fit_with_cov_features_sets_has_cov_flag(self):
        guard = ChatGuard(mode="native")
        guard.fit(LABELED_PAIRS_20, cov_features_list=COV_FEATURES_20)
        assert guard._has_cov_features is True

    def test_fit_without_cov_features_flag_is_false(self):
        guard = ChatGuard(mode="native")
        guard.fit(LABELED_PAIRS_20)
        assert guard._has_cov_features is False

    def test_fit_cov_length_mismatch_raises(self):
        guard = ChatGuard(mode="native")
        with pytest.raises(ValueError, match="length must match"):
            guard.fit(LABELED_PAIRS_20, cov_features_list=COV_FEATURES_20[:5])

    def test_fit_cov_wrong_keys_raises(self):
        guard = ChatGuard(mode="native")
        bad_cov = [{"wrong_key": 0.1}] * 20
        with pytest.raises(ValueError):
            guard.fit(LABELED_PAIRS_20, cov_features_list=bad_cov)

    def test_fit_with_cov_returns_chat_result(self):
        guard = ChatGuard(mode="native")
        guard.fit(LABELED_PAIRS_20, cov_features_list=COV_FEATURES_20)
        result = guard.score_chat("What is Paris?", "Paris is a city.")
        assert isinstance(result, ChatResult)
        assert 0.0 <= result.risk_score <= 1.0


# ── auto_fit() ────────────────────────────────────────────────────────────────

class TestAutoFit:
    def test_auto_fit_raises_on_zero_n_pseudo(self):
        guard = ChatGuard(mode="native")
        with pytest.raises(ValueError):
            guard.auto_fit([{"query": "q", "response": "r"}], n_pseudo=0)

    def test_auto_fit_raises_on_negative_n_pseudo(self):
        guard = ChatGuard(mode="native")
        with pytest.raises(ValueError):
            guard.auto_fit([{"query": "q", "response": "r"}], n_pseudo=-1)

    def test_auto_fit_ignores_label_key(self):
        """auto_fit must ignore any existing label key — no leakage."""
        guard = ChatGuard(mode="native")
        pairs_with_labels = (
            [{"query": f"q{i}", "response": "Paris is the capital.", "label": 99} for i in range(10)]
            + [{"query": f"q{i}", "response": "I think maybe.", "label": 99} for i in range(10)]
        )
        # Alternate YES/NO so fit() has both classes
        call_count = {"n": 0}
        def alternating_create(**kwargs):
            r = MagicMock()
            r.content = [MagicMock(text="YES" if call_count["n"] % 2 == 0 else "NO")]
            call_count["n"] += 1
            return r

        with patch.object(guard, "_get_client") as mock_get_client, \
             patch("llm_guard.chat_guard.time") as mock_time:
            mock_time.sleep = MagicMock()
            mock_anthropic = MagicMock()
            mock_anthropic.messages.create.side_effect = alternating_create
            mock_get_client.return_value = mock_anthropic
            guard.auto_fit(pairs_with_labels, n_pseudo=20)
        assert guard._judge is not None  # fit() was called

    def test_auto_fit_uses_first_n_pseudo_pairs(self):
        """Selection must use first n_pseudo pairs, not random."""
        guard = ChatGuard(mode="native")
        pairs = [{"query": f"q{i}", "response": f"r{i}"} for i in range(100)]
        called_queries = []
        call_count = {"n": 0}

        def fake_create(**kwargs):
            msg = kwargs["messages"][0]["content"]
            called_queries.append(msg)
            r = MagicMock()
            # Alternate YES/NO so fit() has both classes
            r.content = [MagicMock(text="YES" if call_count["n"] % 2 == 0 else "NO")]
            call_count["n"] += 1
            return r

        with patch.object(guard, "_get_client") as mock_get_client, \
             patch("llm_guard.chat_guard.time") as mock_time:
            mock_time.sleep = MagicMock()
            mock_anthropic = MagicMock()
            mock_anthropic.messages.create.side_effect = fake_create
            mock_get_client.return_value = mock_anthropic
            guard.auto_fit(pairs, n_pseudo=5)

        assert len(called_queries) == 5
        assert "r0" in called_queries[0]   # first pair was used
        assert "r4" in called_queries[4]   # fifth pair was used

    def test_auto_fit_warns_if_fewer_pairs_than_n_pseudo(self, caplog):
        guard = ChatGuard(mode="native")
        # Need at least 5 valid labels for auto_fit not to raise, with both classes
        call_count = {"n": 0}
        def alternating_create(**kwargs):
            r = MagicMock()
            r.content = [MagicMock(text="YES" if call_count["n"] % 2 == 0 else "NO")]
            call_count["n"] += 1
            return r

        pairs = [{"query": "q", "response": "r"}] * 6
        with patch.object(guard, "_get_client") as mock_get_client, \
             patch("llm_guard.chat_guard.time") as mock_time:
            mock_time.sleep = MagicMock()
            mock_anthropic = MagicMock()
            mock_anthropic.messages.create.side_effect = alternating_create
            mock_get_client.return_value = mock_anthropic
            with caplog.at_level(logging.WARNING):
                guard.auto_fit(pairs, n_pseudo=50)
        assert any("fewer" in msg.lower() or "6" in msg for msg in caplog.messages)

    def test_auto_fit_raises_if_too_few_valid_labels(self):
        guard = ChatGuard(mode="native")
        pairs = [{"query": "q", "response": "r"}] * 5

        mock_response = MagicMock()
        mock_response.content = [MagicMock(text="MAYBE")]  # all skipped
        with patch.object(guard, "_get_client") as mock_get_client, \
             patch("llm_guard.chat_guard.time") as mock_time:
            mock_time.sleep = MagicMock()
            mock_anthropic = MagicMock()
            mock_anthropic.messages.create.return_value = mock_response
            mock_get_client.return_value = mock_anthropic
            with pytest.raises(ValueError, match="too few valid pseudo-labels"):
                guard.auto_fit(pairs, n_pseudo=5)
