"""
Tests that cache_control appears in Anthropic API calls.
Uses mock patching so no real API key is needed.
"""
from unittest.mock import MagicMock, patch, call as mock_call
import pytest

from llm_guard.agent_guard import AgentGuard


def _make_mock_client(response_text: str = '{"label": "GOOD", "confidence": 0.9}'):
    """Return a mock Anthropic client with a preset text response."""
    mock_client = MagicMock()
    mock_response = MagicMock()
    mock_response.content = [MagicMock(text=response_text)]
    mock_client.messages.create.return_value = mock_response
    return mock_client


STEPS = [
    {"thought": "Let me search", "action_type": "search",
     "action_arg": "capital of France", "observation": "Paris is the capital."}
]


def _assert_system_has_cache_control(mock_client):
    """Assert that at least one messages.create call used cache_control."""
    assert mock_client.messages.create.called, "messages.create was never called"
    for c in mock_client.messages.create.call_args_list:
        kwargs = c.kwargs if c.kwargs else {}
        system_arg = kwargs.get("system")
        if system_arg is None and c.args:
            continue
        if isinstance(system_arg, list):
            for block in system_arg:
                if block.get("cache_control") == {"type": "ephemeral"}:
                    return  # found it
    pytest.fail(
        "No messages.create call had system=[{..., 'cache_control': {'type': 'ephemeral'}}]. "
        f"Actual calls: {mock_client.messages.create.call_args_list}"
    )


class TestRunJudgeCacheControl:
    def test_run_judge_uses_cache_control(self):
        guard = AgentGuard(api_key="fake-key", use_judge=False)
        mock_client = _make_mock_client('{"label": "GOOD", "confidence": 0.9}')
        with patch.object(guard, "_get_client", return_value=mock_client):
            guard._run_judge("What is the capital of France?", STEPS, "Paris")
        _assert_system_has_cache_control(mock_client)

    def test_run_judge_system_is_list_not_str(self):
        guard = AgentGuard(api_key="fake-key")
        mock_client = _make_mock_client('{"label": "GOOD", "confidence": 0.9}')
        with patch.object(guard, "_get_client", return_value=mock_client):
            guard._run_judge("Q?", STEPS, "A")
        call_kwargs = mock_client.messages.create.call_args.kwargs
        assert isinstance(call_kwargs["system"], list), \
            "system should be a list of content blocks for Anthropic cache"


class TestHaikuStepJudgeCacheControl:
    def test_run_haiku_step_judge_uses_cache_control(self):
        guard = AgentGuard(api_key="fake-key")
        mock_client = _make_mock_client(
            '{"step_relevance": 4, "coherence": 4, "on_track": true, "risk_assessment": "LOW"}'
        )
        with patch.object(guard, "_get_client", return_value=mock_client):
            guard._run_haiku_step_judge("Q?", STEPS, k=1)
        _assert_system_has_cache_control(mock_client)


class TestScoreWithPtrueCacheControl:
    def test_score_with_ptrue_uses_cache_control(self):
        guard = AgentGuard(api_key="fake-key")
        mock_client = _make_mock_client("4")
        with patch.object(guard, "_get_client", return_value=mock_client):
            guard.score_with_ptrue("Q?", STEPS, "A")
        _assert_system_has_cache_control(mock_client)


class TestNimPathUnaffected:
    def test_nim_path_receives_string_not_list(self):
        """NIM path must always receive a plain string for system, never a list."""
        guard = AgentGuard(nim_api_key="fake-nim-key")
        mock_nim_client = MagicMock()
        mock_nim_response = MagicMock()
        mock_nim_response.choices = [MagicMock(message=MagicMock(
            content='{"label": "GOOD", "confidence": 0.9}'
        ))]
        mock_nim_client.chat.completions.create.return_value = mock_nim_response
        with patch.object(guard, "_get_nim_client", return_value=mock_nim_client):
            guard._run_judge("Q?", STEPS, "A")
        call_kwargs = mock_nim_client.chat.completions.create.call_args.kwargs
        messages = call_kwargs.get("messages", [])
        system_messages = [m for m in messages if m.get("role") == "system"]
        assert len(system_messages) > 0
        for m in system_messages:
            assert isinstance(m["content"], str), \
                "NIM system content must be a plain string, not a list"
